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

import numpy as np
from typing import Callable, Dict, List, Tuple
 
def monte_carlo_state_value(
    env,
    policy: Callable,
    num_episodes: int = 1000,
    gamma: float = 0.99,
    first_visit: bool = True
) -> Dict[int, float]:
    """
    Monte Carlo estimation of V^π(s) for discrete state spaces.
    
    First-visit MC: Only count the first occurrence of each state per episode.
    Every-visit MC: Count every occurrence of each state.
    
    Args:
        env: Environment with discrete states
        policy: Function mapping state to action
        num_episodes: Number of episodes to sample
        gamma: Discount factor
        first_visit: Use first-visit MC if True, every-visit if False
        
    Returns:
        V: Dictionary mapping states to estimated values
    """
    # Track returns for each state
    returns = {}  # state -> list of returns
    
    for episode in range(num_episodes):
        # Generate episode
        states, rewards = [], []
        state = env.reset()
        done = False
        
        while not done:
            action = policy(state)
            next_state, reward, done, _ = env.step(action)
            states.append(state)
            rewards.append(reward)
            state = next_state
        
        # Compute returns working backwards
        G = 0
        visited = set()
        
        for t in reversed(range(len(states))):
            G = rewards[t] + gamma * G
            s = states[t]
            
            # First-visit: skip if already visited in this episode
            if first_visit and s in visited:
                continue
            visited.add(s)
            
            if s not in returns:
                returns[s] = []
            returns[s].append(G)
    
    # Average returns for each state
    V = {s: np.mean(rs) for s, rs in returns.items()}
    
    return V
 
 
class TDStateValueEstimator:
    """
    Temporal Difference (TD(0)) estimation of V^π(s).
    
    Update rule:
    V(s) ← V(s) + α [r + γV(s') - V(s)]
    
    This is an online algorithm that updates after every step,
    using the bootstrap estimate V(s') instead of waiting for
    the true return.
    """
    
    def __init__(self, n_states: int, 
                 learning_rate: float = 0.1,
                 gamma: float = 0.99):
        self.V = np.zeros(n_states)
        self.alpha = learning_rate
        self.gamma = gamma
    
    def update(self, state: int, reward: float, 
               next_state: int, done: bool) -> float:
        """
        Perform one TD(0) update.
        
        Returns:
            td_error: The temporal difference error δ
        """
        if done:
            # Terminal state has value 0
            target = reward
        else:
            # Bootstrap from next state's value
            target = reward + self.gamma * self.V[next_state]
        
        # TD error: how much our prediction was off
        td_error = target - self.V[state]
        
        # Update value estimate
        self.V[state] += self.alpha * td_error
        
        return td_error
    
    def get_value(self, state: int) -> float:
        return self.V[state]

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

import numpy as np
import torch
from typing import List
 
def compute_advantages_monte_carlo(
    rewards: List[float],
    values: List[float],
    gamma: float = 0.99
) -> np.ndarray:
    """
    Compute advantages using Monte Carlo returns minus value baseline.
    
    A_t = G_t - V(s_t)
    
    High variance but unbiased.
    """
    T = len(rewards)
    returns = np.zeros(T)
    
    # Compute returns backwards
    G = 0
    for t in reversed(range(T)):
        G = rewards[t] + gamma * G
        returns[t] = G
    
    # Advantage = return - baseline
    advantages = returns - np.array(values)
    
    return advantages
 
 
def compute_advantages_td(
    rewards: List[float],
    values: List[float],
    next_values: List[float],
    dones: List[bool],
    gamma: float = 0.99
) -> np.ndarray:
    """
    Compute TD(0) advantages.
    
    A_t = r_t + γV(s_{t+1}) - V(s_t)  (one-step advantage)
    
    Low variance but biased (depends on value function accuracy).
    """
    T = len(rewards)
    advantages = np.zeros(T)
    
    for t in range(T):
        if dones[t]:
            # Terminal: no next state value
            advantages[t] = rewards[t] - values[t]
        else:
            # TD target - current value
            advantages[t] = rewards[t] + gamma * next_values[t] - values[t]
    
    return advantages
 
 
def compute_gae(
    rewards: List[float],
    values: List[float],
    next_values: List[float],
    dones: List[bool],
    gamma: float = 0.99,
    gae_lambda: float = 0.95
) -> np.ndarray:
    """
    Generalized Advantage Estimation (GAE) - Schulman et al., 2016
    
    GAE interpolates between TD (λ=0) and Monte Carlo (λ=1).
    A^GAE = sum_{l=0}^∞ (γλ)^l δ_{t+l}
    
    where δ_t = r_t + γV(s_{t+1}) - V(s_t) is the TD error.
    
    λ controls the bias-variance tradeoff:
    - λ=0: TD(0), low variance, high bias
    - λ=1: Monte Carlo, high variance, low bias
    - λ≈0.95: Common choice, good balance
    
    This is the standard advantage estimator in PPO, A2C, and TRPO.
    """
    T = len(rewards)
    advantages = np.zeros(T)
    
    # Compute backwards for efficiency
    gae = 0
    for t in reversed(range(T)):
        if dones[t]:
            delta = rewards[t] - values[t]
            gae = delta  # Reset GAE at terminal states
        else:
            delta = rewards[t] + gamma * next_values[t] - values[t]
            gae = delta + gamma * gae_lambda * gae
        
        advantages[t] = gae
    
    return advantages
 
 
class GAEEstimator:
    """
    Production-ready GAE estimator with normalization.
    
    Used in PPO and other policy gradient methods.
    """
    
    def __init__(self, gamma: float = 0.99, 
                 gae_lambda: float = 0.95,
                 normalize: bool = True):
        self.gamma = gamma
        self.gae_lambda = gae_lambda
        self.normalize = normalize
    
    def __call__(self,
                 rewards: torch.Tensor,
                 values: torch.Tensor,
                 dones: torch.Tensor) -> torch.Tensor:
        """
        Compute GAE advantages.
        
        Args:
            rewards: [T] tensor of rewards
            values: [T+1] tensor of values (includes bootstrap value)
            dones: [T] tensor of done flags
            
        Returns:
            advantages: [T] tensor of GAE advantages
        """
        T = rewards.shape[0]
        advantages = torch.zeros(T)
        
        gae = 0
        for t in reversed(range(T)):
            # Handle terminal states
            non_terminal = 1.0 - dones[t].float()
            
            # TD error
            delta = rewards[t] + self.gamma * values[t+1] * non_terminal - values[t]
            
            # GAE recursion
            gae = delta + self.gamma * self.gae_lambda * non_terminal * gae
            advantages[t] = gae
        
        # Normalize for stability
        if self.normalize and len(advantages) > 1:
            advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-8)
        
        return advantages

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

import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Tuple
 
class ValueNetwork(nn.Module):
    """
    Neural network for state value function V(s).
    
    Takes state as input, outputs single scalar value.
    """
    
    def __init__(self, state_dim: int, 
                 hidden_dims: Tuple[int, ...] = (64, 64)):
        super().__init__()
        
        layers = []
        prev_dim = state_dim
        
        for hidden_dim in hidden_dims:
            layers.extend([
                nn.Linear(prev_dim, hidden_dim),
                nn.ReLU(),
            ])
            prev_dim = hidden_dim
        
        layers.append(nn.Linear(prev_dim, 1))
        
        self.network = nn.Sequential(*layers)
    
    def forward(self, state: torch.Tensor) -> torch.Tensor:
        """Return value estimate for state(s)."""
        return self.network(state).squeeze(-1)
 
 
class QNetwork(nn.Module):
    """
    Neural network for action value function Q(s,a).
    
    For DISCRETE actions: takes state, outputs Q-value for each action.
    """
    
    def __init__(self, state_dim: int, n_actions: int,
                 hidden_dims: Tuple[int, ...] = (64, 64)):
        super().__init__()
        
        layers = []
        prev_dim = state_dim
        
        for hidden_dim in hidden_dims:
            layers.extend([
                nn.Linear(prev_dim, hidden_dim),
                nn.ReLU(),
            ])
            prev_dim = hidden_dim
        
        layers.append(nn.Linear(prev_dim, n_actions))
        
        self.network = nn.Sequential(*layers)
    
    def forward(self, state: torch.Tensor) -> torch.Tensor:
        """Return Q-values for all actions."""
        return self.network(state)
    
    def get_value(self, state: torch.Tensor, 
                  action: torch.Tensor) -> torch.Tensor:
        """Return Q-value for specific state-action pair."""
        q_values = self.forward(state)
        return q_values.gather(1, action.unsqueeze(1)).squeeze(1)
 
 
class ContinuousQNetwork(nn.Module):
    """
    Q-network for CONTINUOUS actions.
    
    Takes (state, action) pair as input, outputs single Q-value.
    Used in DDPG, TD3, SAC.
    """
    
    def __init__(self, state_dim: int, action_dim: int,
                 hidden_dims: Tuple[int, ...] = (256, 256)):
        super().__init__()
        
        layers = []
        prev_dim = state_dim + action_dim  # Concatenate state and action
        
        for hidden_dim in hidden_dims:
            layers.extend([
                nn.Linear(prev_dim, hidden_dim),
                nn.ReLU(),
            ])
            prev_dim = hidden_dim
        
        layers.append(nn.Linear(prev_dim, 1))
        
        self.network = nn.Sequential(*layers)
    
    def forward(self, state: torch.Tensor, 
                action: torch.Tensor) -> torch.Tensor:
        """Return Q-value for state-action pair."""
        x = torch.cat([state, action], dim=-1)
        return self.network(x).squeeze(-1)
 
 
class DuelingDQN(nn.Module):
    """
    Dueling DQN architecture (Wang et al., 2016).
    
    Separates Q(s,a) into:
    - V(s): state value (how good is this state?)
    - A(s,a): advantage (how much better is this action than average?)
    
    Q(s,a) = V(s) + A(s,a) - mean_a'(A(s,a'))
    
    The subtraction ensures A sums to zero and V represents the true state value.
    This separation helps with credit assignment and learning efficiency.
    """
    
    def __init__(self, state_dim: int, n_actions: int,
                 hidden_dim: int = 128):
        super().__init__()
        
        # Shared feature extraction
        self.features = nn.Sequential(
            nn.Linear(state_dim, hidden_dim),
            nn.ReLU(),
        )
        
        # Value stream: V(s)
        self.value_stream = nn.Sequential(
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, 1)
        )
        
        # Advantage stream: A(s,a)
        self.advantage_stream = nn.Sequential(
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, n_actions)
        )
    
    def forward(self, state: torch.Tensor) -> torch.Tensor:
        features = self.features(state)
        
        value = self.value_stream(features)
        advantages = self.advantage_stream(features)
        
        # Combine: Q = V + (A - mean(A))
        # Subtracting mean ensures identifiability
        q_values = value + (advantages - advantages.mean(dim=-1, keepdim=True))
        
        return q_values

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

import torch
import torch.nn as nn
from copy import deepcopy
 
class TargetNetworkMixin:
    """
    Mixin providing target network functionality.
    
    Used in DQN, DDPG, TD3, SAC, and most deep RL algorithms.
    """
    
    def init_target_networks(self, networks: dict):
        """
        Create target networks as copies of main networks.
        
        Args:
            networks: Dict of {name: network} pairs
        """
        self.target_networks = {}
        for name, network in networks.items():
            # Deep copy to create separate parameters
            self.target_networks[name] = deepcopy(network)
            # Freeze target network (no gradient computation)
            for param in self.target_networks[name].parameters():
                param.requires_grad = False
    
    def soft_update(self, tau: float = 0.005):
        """
        Soft (Polyak) update of target networks.
        
        θ_target = τ × θ + (1-τ) × θ_target
        
        Small τ (0.001-0.01) provides stable, slowly-moving targets.
        """
        for name, target_net in self.target_networks.items():
            main_net = getattr(self, name)
            for target_param, main_param in zip(
                target_net.parameters(), 
                main_net.parameters()
            ):
                target_param.data.copy_(
                    tau * main_param.data + (1 - tau) * target_param.data
                )
    
    def hard_update(self):
        """
        Hard update: copy main network parameters to target.
        
        Used every N steps in original DQN.
        """
        for name, target_net in self.target_networks.items():
            main_net = getattr(self, name)
            target_net.load_state_dict(main_net.state_dict())
 
 
class DQN(nn.Module, TargetNetworkMixin):
    """
    Deep Q-Network with target network.
    
    Key innovations from Mnih et al., 2015:
    1. Experience replay (break correlation)
    2. Target network (stabilize targets)
    """
    
    def __init__(self, state_dim: int, n_actions: int,
                 hidden_dims=(128, 128)):
        super().__init__()
        
        # Main Q-network
        layers = []
        prev = state_dim
        for h in hidden_dims:
            layers.extend([nn.Linear(prev, h), nn.ReLU()])
            prev = h
        layers.append(nn.Linear(prev, n_actions))
        self.q_network = nn.Sequential(*layers)
        
        # Create target network
        self.init_target_networks({'q_network': self.q_network})
        
        self.target_update_freq = 1000
        self.update_count = 0
    
    def get_q_values(self, state: torch.Tensor) -> torch.Tensor:
        """Get Q-values from main network."""
        return self.q_network(state)
    
    def get_target_q_values(self, state: torch.Tensor) -> torch.Tensor:
        """Get Q-values from target network."""
        return self.target_networks['q_network'](state)
    
    def update(self, states, actions, rewards, next_states, dones,
               gamma=0.99):
        """
        Perform one DQN update step.
        """
        # Current Q-values for taken actions
        q_values = self.get_q_values(states).gather(1, actions.unsqueeze(1))
        
        # Target: r + γ max_a' Q_target(s', a')
        with torch.no_grad():
            next_q_values = self.get_target_q_values(next_states)
            max_next_q = next_q_values.max(dim=1)[0]
            targets = rewards + gamma * max_next_q * (1 - dones.float())
        
        # MSE loss
        loss = nn.functional.mse_loss(q_values.squeeze(), targets)
        
        # ... optimizer step ...
        
        # Periodic hard update
        self.update_count += 1
        if self.update_count % self.target_update_freq == 0:
            self.hard_update()
        
        return loss.item()