Deep Reinforcement Learning

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

import torch
import torch.nn as nn
import torch.nn.functional as F
 
class DQN(nn.Module):
    """
    Deep Q-Network Architecture
    
    Processes 84x84x4 frame stacks and outputs Q-values for each action.
    Architecture follows the original DeepMind paper with some modern improvements.
    """
    
    def __init__(self, num_actions: int):
        super(DQN, self).__init__()
        
        # Convolutional feature extractor
        # Designed to progressively extract spatial hierarchies from pixels
        self.conv1 = nn.Conv2d(
            in_channels=4,      # 4 stacked frames
            out_channels=32,    # 32 feature maps
            kernel_size=8,      # Large kernel for initial features
            stride=4            # Aggressive downsampling
        )
        self.conv2 = nn.Conv2d(
            in_channels=32,
            out_channels=64,
            kernel_size=4,
            stride=2
        )
        self.conv3 = nn.Conv2d(
            in_channels=64,
            out_channels=64,
            kernel_size=3,
            stride=1            # Preserve spatial resolution
        )
        
        # Calculate the flattened size after convolutions
        # For 84x84 input: (84-8)/4+1=20, (20-4)/2+1=9, (9-3)/1+1=7
        # Final: 7 * 7 * 64 = 3136
        self.fc_input_size = 7 * 7 * 64
        
        # Fully connected layers for value estimation
        self.fc1 = nn.Linear(self.fc_input_size, 512)
        self.fc2 = nn.Linear(512, num_actions)
        
        # Initialize weights using orthogonal initialization
        self._initialize_weights()
    
    def _initialize_weights(self):
        """
        Proper weight initialization is crucial for training stability.
        Orthogonal initialization helps prevent vanishing/exploding gradients.
        """
        for module in self.modules():
            if isinstance(module, nn.Conv2d):
                nn.init.orthogonal_(module.weight, gain=nn.init.calculate_gain('relu'))
                if module.bias is not None:
                    nn.init.zeros_(module.bias)
            elif isinstance(module, nn.Linear):
                nn.init.orthogonal_(module.weight, gain=1.0)
                if module.bias is not None:
                    nn.init.zeros_(module.bias)
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Forward pass through the network.
        
        Args:
            x: Batch of preprocessed frame stacks, shape (batch, 4, 84, 84)
               Pixel values should be normalized to [0, 1]
        
        Returns:
            Q-values for each action, shape (batch, num_actions)
        """
        # Normalize pixel values if not already done
        if x.max() > 1.0:
            x = x / 255.0
        
        # Convolutional feature extraction
        x = F.relu(self.conv1(x))
        x = F.relu(self.conv2(x))
        x = F.relu(self.conv3(x))
        
        # Flatten spatial dimensions
        x = x.view(x.size(0), -1)  # (batch, 3136)
        
        # Fully connected value estimation
        x = F.relu(self.fc1(x))
        q_values = self.fc2(x)     # No activation - Q-values unrestricted
        
        return q_values
    
    def select_action(self, state: torch.Tensor, epsilon: float = 0.0) -> int:
        """
        Select action using epsilon-greedy policy.
        
        Args:
            state: Single preprocessed frame stack, shape (4, 84, 84)
            epsilon: Probability of random action (exploration)
        
        Returns:
            Selected action index
        """
        if torch.rand(1).item() < epsilon:
            return torch.randint(self.fc2.out_features, (1,)).item()
        
        with torch.no_grad():
            state = state.unsqueeze(0)  # Add batch dimension
            q_values = self.forward(state)
            return q_values.argmax(dim=1).item()

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

import torch
import torch.nn.functional as F
 
def compute_dqn_loss(
    policy_net: DQN,
    target_net: DQN,
    states: torch.Tensor,      # (batch, 4, 84, 84)
    actions: torch.Tensor,     # (batch,) - action indices
    rewards: torch.Tensor,     # (batch,)
    next_states: torch.Tensor, # (batch, 4, 84, 84)
    dones: torch.Tensor,       # (batch,) - episode termination flags
    gamma: float = 0.99
) -> torch.Tensor:
    """
    Compute the DQN loss (Huber loss variant for stability).
    
    The loss measures how well our Q-value predictions match
    the bootstrapped targets from the Bellman equation.
    """
    batch_size = states.size(0)
    
    # Get Q-values for taken actions
    # policy_net outputs shape (batch, num_actions)
    # We select the Q-value corresponding to the action that was actually taken
    q_values = policy_net(states)  # (batch, num_actions)
    q_values = q_values.gather(1, actions.unsqueeze(1)).squeeze(1)  # (batch,)
    
    # Compute target Q-values using target network (no gradients!)
    with torch.no_grad():
        # Get max Q-value for next state from target network
        next_q_values = target_net(next_states)  # (batch, num_actions)
        max_next_q = next_q_values.max(dim=1)[0]  # (batch,)
        
        # Bootstrap target: r + γ * max Q(s', a')
        # If episode terminates (done=True), there's no future value
        targets = rewards + gamma * max_next_q * (1 - dones.float())
    
    # Huber loss (smooth L1) is more robust to outliers than MSE
    # Prevents large gradients from extreme TD errors
    loss = F.smooth_l1_loss(q_values, targets)
    
    return loss
 
# Alternative: Mean Squared Error (original DQN paper)
def compute_dqn_loss_mse(
    policy_net: DQN,
    target_net: DQN,
    states: torch.Tensor,
    actions: torch.Tensor,
    rewards: torch.Tensor,
    next_states: torch.Tensor,
    dones: torch.Tensor,
    gamma: float = 0.99
) -> torch.Tensor:
    """MSE loss variant - less stable but simpler."""
    q_values = policy_net(states).gather(1, actions.unsqueeze(1)).squeeze(1)
    
    with torch.no_grad():
        next_q_values = target_net(next_states).max(dim=1)[0]
        targets = rewards + gamma * next_q_values * (1 - dones.float())
    
    # Mean squared error
    loss = F.mse_loss(q_values, 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
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
207
208
209
210
211
212
213
214
215
216
217

import torch
import torch.optim as optim
from collections import deque
import random
import numpy as np
 
class ReplayBuffer:
    """
    Fixed-size buffer to store experience tuples.
    
    Stores (state, action, reward, next_state, done) tuples and
    supports random sampling for breaking temporal correlations.
    """
    
    def __init__(self, capacity: int = 1_000_000):
        self.buffer = deque(maxlen=capacity)
    
    def push(self, state, action, reward, next_state, done):
        """Store a transition."""
        self.buffer.append((state, action, reward, next_state, done))
    
    def sample(self, batch_size: int):
        """Sample a random batch of transitions."""
        batch = random.sample(self.buffer, batch_size)
        states, actions, rewards, next_states, dones = zip(*batch)
        
        return (
            torch.stack(states),
            torch.tensor(actions, dtype=torch.long),
            torch.tensor(rewards, dtype=torch.float32),
            torch.stack(next_states),
            torch.tensor(dones, dtype=torch.bool)
        )
    
    def __len__(self):
        return len(self.buffer)
 
 
class DQNAgent:
    """
    Complete DQN agent with training logic.
    """
    
    def __init__(
        self,
        num_actions: int,
        device: torch.device = torch.device("cuda"),
        # Hyperparameters (defaults from original paper)
        learning_rate: float = 2.5e-4,
        gamma: float = 0.99,
        buffer_size: int = 1_000_000,
        batch_size: int = 32,
        target_update_freq: int = 10_000,
        epsilon_start: float = 1.0,
        epsilon_end: float = 0.1,
        epsilon_decay_steps: int = 1_000_000,
        learning_starts: int = 50_000,
    ):
        self.device = device
        self.num_actions = num_actions
        self.gamma = gamma
        self.batch_size = batch_size
        self.target_update_freq = target_update_freq
        self.learning_starts = learning_starts
        
        # Epsilon schedule for exploration
        self.epsilon_start = epsilon_start
        self.epsilon_end = epsilon_end
        self.epsilon_decay_steps = epsilon_decay_steps
        
        # Networks
        self.policy_net = DQN(num_actions).to(device)
        self.target_net = DQN(num_actions).to(device)
        self.target_net.load_state_dict(self.policy_net.state_dict())
        self.target_net.eval()  # Target net doesn't need gradients
        
        # Optimizer and replay buffer
        self.optimizer = optim.Adam(
            self.policy_net.parameters(),
            lr=learning_rate
        )
        self.replay_buffer = ReplayBuffer(buffer_size)
        
        # Tracking
        self.total_steps = 0
    
    def get_epsilon(self) -> float:
        """Linear epsilon decay schedule."""
        progress = min(self.total_steps / self.epsilon_decay_steps, 1.0)
        return self.epsilon_start + progress * (self.epsilon_end - self.epsilon_start)
    
    def select_action(self, state: torch.Tensor) -> int:
        """Epsilon-greedy action selection."""
        epsilon = self.get_epsilon()
        
        if random.random() < epsilon:
            return random.randrange(self.num_actions)
        
        with torch.no_grad():
            state = state.unsqueeze(0).to(self.device)
            q_values = self.policy_net(state)
            return q_values.argmax(dim=1).item()
    
    def train_step(self) -> float:
        """Perform one gradient update."""
        if len(self.replay_buffer) < self.batch_size:
            return 0.0
        
        # Sample batch
        states, actions, rewards, next_states, dones = self.replay_buffer.sample(
            self.batch_size
        )
        
        # Move to device
        states = states.to(self.device)
        actions = actions.to(self.device)
        rewards = rewards.to(self.device)
        next_states = next_states.to(self.device)
        dones = dones.to(self.device)
        
        # Compute loss
        loss = compute_dqn_loss(
            self.policy_net,
            self.target_net,
            states, actions, rewards, next_states, dones,
            gamma=self.gamma
        )
        
        # Optimize
        self.optimizer.zero_grad()
        loss.backward()
        # Gradient clipping for stability
        torch.nn.utils.clip_grad_norm_(self.policy_net.parameters(), max_norm=10)
        self.optimizer.step()
        
        return loss.item()
    
    def update_target_network(self):
        """Hard update: copy policy network parameters to target network."""
        self.target_net.load_state_dict(self.policy_net.state_dict())
    
    def train_episode(self, env) -> dict:
        """
        Train for one episode.
        
        Returns:
            Dictionary with episode statistics
        """
        state = env.reset()
        state = preprocess_frame(state)  # Convert to tensor
        frame_stack = deque([state] * 4, maxlen=4)
        
        episode_reward = 0
        episode_loss = 0
        episode_steps = 0
        
        done = False
        while not done:
            # Stack frames for temporal information
            stacked_state = torch.stack(list(frame_stack))
            
            # Select and execute action
            action = self.select_action(stacked_state)
            next_state, reward, done, info = env.step(action)
            
            # Preprocess and update frame stack
            next_state = preprocess_frame(next_state)
            frame_stack.append(next_state)
            next_stacked = torch.stack(list(frame_stack))
            
            # Store transition
            self.replay_buffer.push(
                stacked_state, action, reward, next_stacked, done
            )
            
            # Train if enough samples and past initial exploration
            if self.total_steps >= self.learning_starts:
                loss = self.train_step()
                episode_loss += loss
                
                # Update target network periodically
                if self.total_steps % self.target_update_freq == 0:
                    self.update_target_network()
            
            episode_reward += reward
            episode_steps += 1
            self.total_steps += 1
        
        return {
            "reward": episode_reward,
            "steps": episode_steps,
            "loss": episode_loss / max(episode_steps, 1),
            "epsilon": self.get_epsilon(),
            "total_steps": self.total_steps,
        }
 
 
def preprocess_frame(frame: np.ndarray) -> torch.Tensor:
    """
    Convert raw Atari frame to DQN input format.
    
    - Convert to grayscale
    - Resize to 84x84
    - Convert to tensor and normalize
    """
    import cv2
    
    # Grayscale conversion
    gray = cv2.cvtColor(frame, cv2.COLOR_RGB2GRAY)
    
    # Resize to 84x84
    resized = cv2.resize(gray, (84, 84), interpolation=cv2.INTER_AREA)
    
    # Convert to tensor and normalize to [0, 1]
    tensor = torch.tensor(resized, dtype=torch.float32) / 255.0
    
    return tensor

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

import numpy as np
import torch
 
class EfficientReplayBuffer:
    """
    Memory-efficient replay buffer using pre-allocated numpy arrays.
    
    Key optimizations:
    - Pre-allocated memory prevents fragmentation
    - uint8 storage for observations reduces memory 4x
    - Lazy frame stacking constructs stacks only when sampling
    - Circular buffer with pointer arithmetic
    """
    
    def __init__(
        self,
        capacity: int = 1_000_000,
        frame_shape: tuple = (84, 84),
        frame_stack: int = 4,
    ):
        self.capacity = capacity
        self.frame_shape = frame_shape
        self.frame_stack = frame_stack
        self.ptr = 0      # Current write position
        self.size = 0     # Current buffer size
        
        # Pre-allocate arrays
        # Use uint8 for observations (0-255) to save memory
        self.observations = np.zeros(
            (capacity, *frame_shape), 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_)
    
    def push(self, obs: np.ndarray, action: int, reward: float, done: bool):
        """
        Store a single frame transition.
        
        Note: We store individual frames, not stacked frames.
        This dramatically reduces memory usage.
        """
        self.observations[self.ptr] = obs
        self.actions[self.ptr] = action
        self.rewards[self.ptr] = reward
        self.dones[self.ptr] = done
        
        self.ptr = (self.ptr + 1) % self.capacity
        self.size = min(self.size + 1, self.capacity)
    
    def _get_frame_stack(self, idx: int) -> np.ndarray:
        """
        Construct a frame stack ending at idx.
        
        Handles episode boundaries: if a 'done' flag is encountered
        while looking back, we repeat the frame after 'done'.
        """
        frames = []
        for i in range(self.frame_stack):
            frame_idx = (idx - i) % self.capacity
            
            # Check if we've crossed an episode boundary
            if i > 0 and self.dones[(idx - i + 1) % self.capacity]:
                # Repeat the first frame of the new episode
                frames.append(self.observations[(idx - i + 1) % self.capacity])
            else:
                frames.append(self.observations[frame_idx])
        
        # Reverse to get chronological order
        return np.stack(frames[::-1])
    
    def sample(self, batch_size: int) -> tuple:
        """
        Sample a batch with lazy frame stack construction.
        
        Returns tensors ready for GPU training.
        """
        # Sample indices (avoid the most recent frame_stack-1 entries)
        valid_indices = np.arange(self.frame_stack - 1, self.size)
        indices = np.random.choice(valid_indices, size=batch_size, replace=False)
        
        # Construct frame stacks
        states = np.array([self._get_frame_stack(i) for i in indices])
        next_states = np.array([self._get_frame_stack((i + 1) % self.capacity) for i in indices])
        
        # Convert to tensors and normalize
        return (
            torch.tensor(states, dtype=torch.float32) / 255.0,
            torch.tensor(self.actions[indices], dtype=torch.long),
            torch.tensor(self.rewards[indices], dtype=torch.float32),
            torch.tensor(next_states, dtype=torch.float32) / 255.0,
            torch.tensor(self.dones[indices], dtype=torch.bool),
        )
    
    def __len__(self):
        return self.size
 
 
# Memory comparison
def memory_comparison():
    """
    Naive approach: Store (4, 84, 84) float32 states
    Memory per transition: 4 * 84 * 84 * 4 bytes = 112,896 bytes
    1M transitions: ~108 GB
    
    Efficient approach: Store (84, 84) uint8 observations
    Memory per transition: 84 * 84 * 1 byte = 7,056 bytes
    1M transitions: ~6.7 GB
    
    Savings: 16x memory reduction!
    """
    pass