1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119

import torch
import torch.nn as nn
import torch.nn.functional as F
 
class DuelingDQN(nn.Module):
    """
    Dueling DQN architecture that separates value and advantage streams.
    
    The key insight: many states have similar values regardless of action.
    By explicitly modeling state value, we can generalize better.
    """
    
    def __init__(self, num_actions: int, num_atoms: int = 1):
        """
        Args:
            num_actions: Number of discrete actions
            num_atoms: Number of atoms for distributional RL (1 for standard)
        """
        super(DuelingDQN, self).__init__()
        
        self.num_actions = num_actions
        self.num_atoms = num_atoms
        
        # Shared convolutional backbone
        self.conv1 = nn.Conv2d(4, 32, kernel_size=8, stride=4)
        self.conv2 = nn.Conv2d(32, 64, kernel_size=4, stride=2)
        self.conv3 = nn.Conv2d(64, 64, kernel_size=3, stride=1)
        
        # Calculate flattened size: 7 * 7 * 64 = 3136
        self.feature_size = 7 * 7 * 64
        
        # VALUE STREAM: estimates V(s)
        # Outputs a single value (or distribution if num_atoms > 1)
        self.value_stream = nn.Sequential(
            nn.Linear(self.feature_size, 512),
            nn.ReLU(),
            nn.Linear(512, num_atoms)  # Single value per atom
        )
        
        # ADVANTAGE STREAM: estimates A(s, a)
        # Outputs advantages for each action
        self.advantage_stream = nn.Sequential(
            nn.Linear(self.feature_size, 512),
            nn.ReLU(),
            nn.Linear(512, num_actions * num_atoms)
        )
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Forward pass with dueling aggregation.
        
        Returns Q(s,a) = V(s) + (A(s,a) - mean(A(s,:)))
        """
        batch_size = x.size(0)
        
        # Normalize input
        if x.max() > 1.0:
            x = x / 255.0
        
        # Shared convolutional features
        x = F.relu(self.conv1(x))
        x = F.relu(self.conv2(x))
        x = F.relu(self.conv3(x))
        x = x.view(batch_size, -1)  # Flatten
        
        # Compute value and advantages separately
        value = self.value_stream(x)  # (batch, num_atoms)
        advantages = self.advantage_stream(x)  # (batch, num_actions * num_atoms)
        
        if self.num_atoms > 1:
            # Distributional case: reshape for proper broadcasting
            value = value.view(batch_size, 1, self.num_atoms)
            advantages = advantages.view(batch_size, self.num_actions, self.num_atoms)
            
            # Combine with mean-centered advantages
            # Q = V + (A - mean(A))
            q_values = value + (advantages - advantages.mean(dim=1, keepdim=True))
            return q_values  # (batch, num_actions, num_atoms)
        
        else:
            # Standard case
            value = value.view(batch_size, 1)
            advantages = advantages.view(batch_size, self.num_actions)
            
            # Combine: Q(s,a) = V(s) + A(s,a) - mean_a(A(s,a))
            q_values = value + (advantages - advantages.mean(dim=1, keepdim=True))
            return q_values  # (batch, num_actions)
 
 
def visualize_dueling_benefit():
    """
    Demonstrate when dueling architecture helps.
    
    When actions have similar values, dueling learns faster
    because V(s) is shared across all actions.
    """
    import numpy as np
    
    # Scenario: 4 actions, but only one matters
    # True Q-values:
    # Q(s, a0) = 1.0
    # Q(s, a1) = 1.01  (slightly better)
    # Q(s, a2) = 0.99
    # Q(s, a3) = 1.0
    
    true_v = 1.0  # State value
    true_advantages = np.array([0.0, 0.01, -0.01, 0.0])  # Small differences
    
    # Standard DQN: must learn all 4 Q-values independently
    # Dueling DQN: learns V once, then small advantages
    
    # Number of samples to learn (simplified)
    samples_standard = 4  # Need to visit each action
    samples_dueling = 1 + 0.5  # V from any action + subtle differences
    
    print("Efficiency comparison:")
    print(f"Standard DQN: ~{samples_standard}x samples per state")
    print(f"Dueling DQN: ~{samples_dueling}x samples per state")
    print(f"Advantage: {samples_standard / samples_dueling:.1f}x faster")

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128

import torch
import numpy as np
from collections import deque
 
class MultiStepBuffer:
    """
    Experience replay buffer that computes n-step returns.
    
    Stores individual transitions but yields n-step returns for training.
    """
    
    def __init__(
        self,
        capacity: int = 1_000_000,
        n_step: int = 3,
        gamma: float = 0.99
    ):
        self.capacity = capacity
        self.n_step = n_step
        self.gamma = gamma
        
        # Main storage
        self.states = np.zeros((capacity, 4, 84, 84), dtype=np.uint8)
        self.actions = np.zeros(capacity, dtype=np.int32)
        self.rewards = np.zeros(capacity, dtype=np.float32)
        self.dones = np.zeros(capacity, dtype=np.bool_)
        
        # For n-step computation
        self.n_step_buffer = deque(maxlen=n_step)
        
        self.ptr = 0
        self.size = 0
    
    def push(self, state, action, reward, done):
        """
        Add a transition to the n-step buffer.
        
        When buffer is full, compute n-step return and store.
        """
        self.n_step_buffer.append((state, action, reward, done))
        
        # Wait until we have n transitions
        if len(self.n_step_buffer) < self.n_step:
            return
        
        # Compute n-step return
        n_step_return = 0.0
        for i, (s, a, r, d) in enumerate(self.n_step_buffer):
            n_step_return += (self.gamma ** i) * r
            if d:  # Episode terminated
                break
        
        # Store the first transition with n-step return
        first_state, first_action, _, _ = self.n_step_buffer[0]
        
        self.states[self.ptr] = first_state
        self.actions[self.ptr] = first_action
        self.rewards[self.ptr] = n_step_return  # N-step discounted return
        self.dones[self.ptr] = done or any(t[3] for t in self.n_step_buffer)
        
        self.ptr = (self.ptr + 1) % self.capacity
        self.size = min(self.size + 1, self.capacity)
    
    def finish_episode(self):
        """
        Flush remaining transitions at episode end.
        
        Must be called when episode terminates to store remaining n-step returns.
        """
        while len(self.n_step_buffer) > 0:
            n_step_return = 0.0
            for i, (s, a, r, d) in enumerate(self.n_step_buffer):
                n_step_return += (self.gamma ** i) * r
                if d:
                    break
            
            first_state, first_action, _, _ = self.n_step_buffer[0]
            
            self.states[self.ptr] = first_state
            self.actions[self.ptr] = first_action
            self.rewards[self.ptr] = n_step_return
            self.dones[self.ptr] = True  # Episode ended
            
            self.ptr = (self.ptr + 1) % self.capacity
            self.size = min(self.size + 1, self.capacity)
            
            self.n_step_buffer.popleft()
    
    def sample(self, batch_size: int, device: torch.device):
        """
        Sample batch with n-step returns already computed.
        """
        indices = np.random.randint(0, self.size, size=batch_size)
        
        # Get states n steps ahead for bootstrapping
        next_indices = (indices + self.n_step) % self.capacity
        next_indices = np.clip(next_indices, 0, self.size - 1)
        
        return {
            'states': torch.FloatTensor(self.states[indices]).to(device) / 255.0,
            'actions': torch.LongTensor(self.actions[indices]).to(device),
            'rewards': torch.FloatTensor(self.rewards[indices]).to(device),
            'next_states': torch.FloatTensor(self.states[next_indices]).to(device) / 255.0,
            'dones': torch.BoolTensor(self.dones[indices]).to(device),
            'gamma_n': self.gamma ** self.n_step,  # Discount for n steps
        }
 
 
def compute_nstep_loss(batch, policy_net, target_net, gamma_n):
    """
    Compute loss using n-step returns.
    
    The key difference: we use gamma^n for bootstrapping, not gamma.
    """
    # Current Q-values
    current_q = policy_net(batch['states'])
    current_q = current_q.gather(1, batch['actions'].unsqueeze(1)).squeeze(1)
    
    with torch.no_grad():
        # Double DQN: policy selects, target evaluates
        next_actions = policy_net(batch['next_states']).argmax(dim=1, keepdim=True)
        next_q = target_net(batch['next_states']).gather(1, next_actions).squeeze(1)
        
        # N-step target: sum of n rewards + gamma^n * Q(s_{t+n})
        targets = batch['rewards'] + gamma_n * next_q * (~batch['dones']).float()
    
    loss = F.smooth_l1_loss(current_q, targets)
    return loss

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162

import torch
import torch.nn as nn
import torch.nn.functional as F
 
class DistributionalDQN(nn.Module):
    """
    Distributional DQN (C51) with categorical distribution over returns.
    
    Instead of outputting Q(s,a), outputs the distribution of returns.
    The Q-value is the expectation under this distribution.
    """
    
    def __init__(
        self,
        num_actions: int,
        num_atoms: int = 51,
        v_min: float = -10.0,
        v_max: float = 10.0
    ):
        super().__init__()
        
        self.num_actions = num_actions
        self.num_atoms = num_atoms
        self.v_min = v_min
        self.v_max = v_max
        
        # Support: fixed atoms where probabilities are defined
        self.register_buffer(
            'support',
            torch.linspace(v_min, v_max, num_atoms)
        )
        self.delta_z = (v_max - v_min) / (num_atoms - 1)
        
        # Convolutional backbone
        self.conv = nn.Sequential(
            nn.Conv2d(4, 32, kernel_size=8, stride=4),
            nn.ReLU(),
            nn.Conv2d(32, 64, kernel_size=4, stride=2),
            nn.ReLU(),
            nn.Conv2d(64, 64, kernel_size=3, stride=1),
            nn.ReLU(),
        )
        
        # Output: num_actions * num_atoms logits
        self.fc = nn.Sequential(
            nn.Linear(7 * 7 * 64, 512),
            nn.ReLU(),
            nn.Linear(512, num_actions * num_atoms)  # Distribution per action
        )
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Returns probability distributions over returns for each action.
        
        Output shape: (batch, num_actions, num_atoms)
        """
        if x.max() > 1.0:
            x = x / 255.0
        
        batch_size = x.size(0)
        features = self.conv(x)
        features = features.view(batch_size, -1)
        
        # Raw logits
        logits = self.fc(features)
        logits = logits.view(batch_size, self.num_actions, self.num_atoms)
        
        # Convert to probabilities (softmax over atoms)
        probs = F.softmax(logits, dim=-1)
        
        return probs
    
    def get_q_values(self, x: torch.Tensor) -> torch.Tensor:
        """
        Compute Q-values as expectation of return distribution.
        
        Q(s,a) = sum_i p_i(s,a) * z_i
        """
        probs = self.forward(x)  # (batch, num_actions, num_atoms)
        q_values = (probs * self.support).sum(dim=-1)  # (batch, num_actions)
        return q_values
 
 
def compute_distributional_loss(
    policy_net: DistributionalDQN,
    target_net: DistributionalDQN,
    states: torch.Tensor,
    actions: torch.Tensor,
    rewards: torch.Tensor,
    next_states: torch.Tensor,
    dones: torch.Tensor,
    gamma: float = 0.99
) -> torch.Tensor:
    """
    Compute distributional (C51) loss using categorical cross-entropy.
    
    The key is projecting the target distribution onto the fixed support.
    """
    batch_size = states.size(0)
    support = policy_net.support
    num_atoms = policy_net.num_atoms
    v_min = policy_net.v_min
    v_max = policy_net.v_max
    delta_z = policy_net.delta_z
    
    # Get current distributions
    current_dist = policy_net(states)  # (batch, num_actions, num_atoms)
    # Select distribution for taken actions
    current_dist = current_dist[
        torch.arange(batch_size), actions
    ]  # (batch, num_atoms)
    
    with torch.no_grad():
        # Double DQN: use policy net to select actions
        next_q = policy_net.get_q_values(next_states)
        next_actions = next_q.argmax(dim=1)  # (batch,)
        
        # Get target distribution for those actions
        next_dist = target_net(next_states)  # (batch, num_actions, num_atoms)
        next_dist = next_dist[
            torch.arange(batch_size), next_actions
        ]  # (batch, num_atoms)
        
        # PROJECT target distribution onto support
        # This is the key operation in distributional RL
        
        # Compute target support: r + gamma * z_i (clamped to [v_min, v_max])
        target_support = rewards.unsqueeze(1) + gamma * support * (1 - dones.unsqueeze(1).float())
        target_support = target_support.clamp(v_min, v_max)
        
        # Compute projection indices and weights
        # Each target atom distributes its probability to neighboring fixed atoms
        b = (target_support - v_min) / delta_z  # Fractional indices
        l = b.floor().long()  # Lower index
        u = b.ceil().long()   # Upper index
        
        # Handle edge cases
        l = l.clamp(0, num_atoms - 1)
        u = u.clamp(0, num_atoms - 1)
        
        # Initialize target distribution
        target_dist = torch.zeros_like(next_dist)
        
        # Distribute probability to lower and upper neighbors
        offset = torch.arange(batch_size).unsqueeze(1).expand(batch_size, num_atoms)
        
        target_dist.view(-1).index_add_(
            0,
            (offset * num_atoms + l).view(-1),
            (next_dist * (u.float() - b)).view(-1)
        )
        target_dist.view(-1).index_add_(
            0,
            (offset * num_atoms + u).view(-1),
            (next_dist * (b - l.float())).view(-1)
        )
    
    # Cross-entropy loss between projected target and predicted distribution
    log_probs = torch.log(current_dist + 1e-8)
    loss = -(target_dist * log_probs).sum(dim=-1).mean()
    
    return loss

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167

import torch
import torch.nn as nn
import torch.nn.functional as F
import math
 
class NoisyLinear(nn.Module):
    """
    Noisy Linear layer for exploration in Deep RL.
    
    Implements factorized Gaussian noise for efficiency.
    The network learns the noise scale, enabling state-dependent exploration.
    """
    
    def __init__(
        self,
        in_features: int,
        out_features: int,
        sigma_init: float = 0.5
    ):
        super().__init__()
        
        self.in_features = in_features
        self.out_features = out_features
        self.sigma_init = sigma_init
        
        # Learnable parameters (mean weights and biases)
        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)
        )
        
        # Noise buffers (not learnable, re-sampled each forward pass)
        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 reset_parameters(self):
        """Initialize parameters following the paper."""
        bound = 1 / math.sqrt(self.in_features)
        
        # Mean weights: uniform initialization
        self.weight_mu.data.uniform_(-bound, bound)
        self.bias_mu.data.uniform_(-bound, bound)
        
        # Noise scale: constant initialization
        self.weight_sigma.data.fill_(
            self.sigma_init / math.sqrt(self.in_features)
        )
        self.bias_sigma.data.fill_(
            self.sigma_init / math.sqrt(self.out_features)
        )
    
    def reset_noise(self):
        """
        Re-sample noise for next forward pass.
        
        Uses factorized noise for efficiency:
        - Full noise: O(in * out) random samples
        - Factorized: O(in + out) random samples
        """
        # Factorized Gaussian noise
        epsilon_in = self._scale_noise(self.in_features)
        epsilon_out = self._scale_noise(self.out_features)
        
        # Outer product for weight noise
        self.weight_epsilon.copy_(epsilon_out.outer(epsilon_in))
        self.bias_epsilon.copy_(epsilon_out)
    
    def _scale_noise(self, size: int) -> torch.Tensor:
        """
        Generate scaled noise using the paper's transformation.
        
        f(x) = sign(x) * sqrt(|x|) makes tails heavier than Gaussian.
        """
        x = torch.randn(size, device=self.weight_mu.device)
        return x.sign() * x.abs().sqrt()
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Forward pass with noisy weights.
        
        If in training mode, uses noisy weights.
        If in eval mode, uses mean weights only (deterministic).
        """
        if self.training:
            # Noisy weights: μ + σ * ε
            weight = self.weight_mu + self.weight_sigma * self.weight_epsilon
            bias = self.bias_mu + self.bias_sigma * self.bias_epsilon
        else:
            # Deterministic: use means only
            weight = self.weight_mu
            bias = self.bias_mu
        
        return F.linear(x, weight, bias)
 
 
class NoisyDuelingDQN(nn.Module):
    """
    Rainbow-style network combining dueling architecture with noisy layers.
    """
    
    def __init__(self, num_actions: int, num_atoms: int = 51):
        super().__init__()
        
        self.num_actions = num_actions
        self.num_atoms = num_atoms
        
        # Shared convolutional backbone (no noise here)
        self.conv = nn.Sequential(
            nn.Conv2d(4, 32, kernel_size=8, stride=4),
            nn.ReLU(),
            nn.Conv2d(32, 64, kernel_size=4, stride=2),
            nn.ReLU(),
            nn.Conv2d(64, 64, kernel_size=3, stride=1),
            nn.ReLU(),
        )
        
        # Dueling streams with noisy layers
        # Value stream
        self.value_fc1 = NoisyLinear(7 * 7 * 64, 512)
        self.value_fc2 = NoisyLinear(512, num_atoms)
        
        # Advantage stream
        self.advantage_fc1 = NoisyLinear(7 * 7 * 64, 512)
        self.advantage_fc2 = NoisyLinear(512, num_actions * num_atoms)
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if x.max() > 1.0:
            x = x / 255.0
        
        batch_size = x.size(0)
        
        # Convolutional features
        features = self.conv(x)
        features = features.view(batch_size, -1)
        
        # Value stream
        value = F.relu(self.value_fc1(features))
        value = self.value_fc2(value)
        value = value.view(batch_size, 1, self.num_atoms)
        
        # Advantage stream
        advantage = F.relu(self.advantage_fc1(features))
        advantage = self.advantage_fc2(advantage)
        advantage = advantage.view(batch_size, self.num_actions, self.num_atoms)
        
        # Combine (dueling aggregation)
        q_dist = value + advantage - advantage.mean(dim=1, keepdim=True)
        
        # Softmax over atoms to get probabilities
        return F.softmax(q_dist, dim=-1)
    
    def reset_noise(self):
        """Reset noise in all noisy layers."""
        for module in self.modules():
            if isinstance(module, NoisyLinear):
                module.reset_noise()

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206

import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np
from typing import Dict, Tuple
 
class RainbowAgent:
    """
    Complete Rainbow DQN implementation.
    
    Combines all six improvements into a single, coherent agent.
    """
    
    def __init__(
        self,
        num_actions: int,
        device: torch.device,
        # Network parameters
        num_atoms: int = 51,
        v_min: float = -10.0,
        v_max: float = 10.0,
        # Multi-step
        n_step: int = 3,
        gamma: float = 0.99,
        # Prioritized replay
        buffer_size: int = 1_000_000,
        alpha: float = 0.6,
        beta_start: float = 0.4,
        beta_frames: int = 100_000,
        # Training
        batch_size: int = 32,
        learning_rate: float = 6.25e-5,
        adam_epsilon: float = 1.5e-4,
        target_update_freq: int = 8000,
    ):
        self.device = device
        self.num_actions = num_actions
        self.gamma = gamma
        self.gamma_n = gamma ** n_step
        self.n_step = n_step
        self.batch_size = batch_size
        self.target_update_freq = target_update_freq
        
        # Distributional parameters
        self.num_atoms = num_atoms
        self.v_min = v_min
        self.v_max = v_max
        self.support = torch.linspace(v_min, v_max, num_atoms).to(device)
        self.delta_z = (v_max - v_min) / (num_atoms - 1)
        
        # Networks (Dueling + Noisy + Distributional)
        self.policy_net = NoisyDuelingDQN(num_actions, num_atoms).to(device)
        self.target_net = NoisyDuelingDQN(num_actions, num_atoms).to(device)
        self.target_net.load_state_dict(self.policy_net.state_dict())
        self.target_net.eval()
        
        # Optimizer (specific settings from paper)
        self.optimizer = optim.Adam(
            self.policy_net.parameters(),
            lr=learning_rate,
            eps=adam_epsilon
        )
        
        # Prioritized replay buffer with n-step support
        self.replay_buffer = PrioritizedMultiStepBuffer(
            capacity=buffer_size,
            n_step=n_step,
            gamma=gamma,
            alpha=alpha,
            beta_start=beta_start,
            beta_frames=beta_frames
        )
        
        self.step_count = 0
    
    def select_action(self, state: np.ndarray) -> int:
        """
        Select action using noisy network (no epsilon-greedy needed).
        """
        self.policy_net.eval()
        
        with torch.no_grad():
            state_tensor = torch.FloatTensor(state).unsqueeze(0).to(self.device)
            if state_tensor.max() > 1.0:
                state_tensor = state_tensor / 255.0
            
            # Reset noise for this action selection
            self.policy_net.reset_noise()
            
            # Get distribution and compute Q-values
            dist = self.policy_net(state_tensor)  # (1, num_actions, num_atoms)
            q_values = (dist * self.support).sum(dim=-1)  # (1, num_actions)
            action = q_values.argmax(dim=-1).item()
        
        self.policy_net.train()
        return action
    
    def train_step(self) -> Dict[str, float]:
        """
        Perform one training step using all Rainbow components.
        """
        if len(self.replay_buffer) < self.batch_size:
            return {}
        
        self.step_count += 1
        
        # PRIORITIZED REPLAY: Sample with priorities
        batch, indices, weights = self.replay_buffer.sample(
            self.batch_size, self.device
        )
        
        # Reset noise for training
        self.policy_net.reset_noise()
        self.target_net.reset_noise()
        
        # Compute loss and TD errors
        loss, td_errors = self._compute_loss(batch, weights)
        
        # Optimize
        self.optimizer.zero_grad()
        loss.backward()
        torch.nn.utils.clip_grad_norm_(self.policy_net.parameters(), 10)
        self.optimizer.step()
        
        # PRIORITIZED REPLAY: Update priorities with new TD errors
        self.replay_buffer.update_priorities(indices, td_errors.detach().cpu().numpy())
        
        # Target network update
        if self.step_count % self.target_update_freq == 0:
            self.target_net.load_state_dict(self.policy_net.state_dict())
        
        return {
            'loss': loss.item(),
            'mean_q': (batch['q_values'] * self.support).sum(-1).mean().item(),
            'mean_priority': weights.mean().item(),
        }
    
    def _compute_loss(
        self,
        batch: Dict[str, torch.Tensor],
        weights: torch.Tensor
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        Compute Rainbow loss combining all components.
        """
        states = batch['states']
        actions = batch['actions']
        rewards = batch['rewards']  # N-step returns
        next_states = batch['next_states']
        dones = batch['dones']
        
        batch_size = states.size(0)
        
        # Current distribution for taken actions
        current_dist = self.policy_net(states)  # (batch, actions, atoms)
        current_dist = current_dist[
            torch.arange(batch_size), actions
        ]  # (batch, atoms)
        
        with torch.no_grad():
            # DOUBLE DQN: Policy net selects action
            next_dist = self.policy_net(next_states)
            next_q = (next_dist * self.support).sum(dim=-1)
            next_actions = next_q.argmax(dim=1)
            
            # Target net evaluates that action
            target_dist = self.target_net(next_states)
            target_dist = target_dist[
                torch.arange(batch_size), next_actions
            ]
            
            # Project target distribution (DISTRIBUTIONAL RL)
            # MULTI-STEP: Use gamma^n
            target_support = rewards.unsqueeze(1) +                             self.gamma_n * self.support * (1 - dones.unsqueeze(1).float())
            target_support = target_support.clamp(self.v_min, self.v_max)
            
            # Compute projection
            b = (target_support - self.v_min) / self.delta_z
            l = b.floor().long().clamp(0, self.num_atoms - 1)
            u = b.ceil().long().clamp(0, self.num_atoms - 1)
            
            projected_dist = torch.zeros_like(target_dist)
            offset = torch.arange(batch_size).unsqueeze(1).expand_as(l).to(self.device)
            
            projected_dist.view(-1).index_add_(
                0,
                (offset * self.num_atoms + l).view(-1),
                (target_dist * (u.float() - b)).view(-1)
            )
            projected_dist.view(-1).index_add_(
                0,
                (offset * self.num_atoms + u).view(-1),
                (target_dist * (b - l.float())).view(-1)
            )
        
        # Cross-entropy loss
        log_probs = torch.log(current_dist + 1e-8)
        elementwise_loss = -(projected_dist * log_probs).sum(dim=-1)
        
        # TD errors for priority updates
        td_errors = elementwise_loss.detach()
        
        # PRIORITIZED REPLAY: Weight the loss
        loss = (elementwise_loss * weights).mean()
        
        return loss, td_errors