1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17

# The advantage measures relative action quality
 
A^π(s, a) = Q^π(s, a) - V^π(s)
 
           = Q^π(s, a) - E_{a'~π}[Q^π(s, a')]
 
           = "How much better is action a than average?"
 
# Properties:
# - A^π(s, a) > 0: Action a is better than average in state s
# - A^π(s, a) = 0: Action a is exactly average
# - A^π(s, a) < 0: Action a is worse than average
 
# Critical insight: The average advantage is always zero!
E_{a~π}[A^π(s, a)] = E_{a~π}[Q^π(s, a)] - V^π(s)
                    = V^π(s) - V^π(s)
                    = 0

1
2
3
4
5
6
7
8
9
10
11
12

# The advantage relates to TD errors
 
# TD error (temporal difference):
δ^π(s, a, s') = r(s, a) + γV^π(s') - V^π(s)
 
# Expectation over next states:
E_{s'|s,a}[δ^π(s, a, s')] = E_{s'|s,a}[r + γV^π(s')] - V^π(s)
                          = Q^π(s, a) - V^π(s)
                          = A^π(s, a)
 
# Key insight: The expected TD error IS the advantage!
# This connects advantages to bootstrapped value estimates

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16

# The policy gradient can be written in multiple equivalent forms:
 
# Form 1: With Q-function
∇_θ J(θ) = E_{s~ρ^π, a~π}[∇_θ log π_θ(a|s) · Q^π(s, a)]
 
# Form 2: With advantage function
∇_θ J(θ) = E_{s~ρ^π, a~π}[∇_θ log π_θ(a|s) · A^π(s, a)]
 
# These are equivalent because:
E_{a~π}[∇_θ log π_θ(a|s) · V^π(s)]
    = V^π(s) · E_{a~π}[∇_θ log π_θ(a|s)]
    = V^π(s) · 0  (score has zero mean)
    = 0
 
# So subtracting V^π(s) from Q^π(s,a) doesn't change the gradient!
# But it DOES reduce variance dramatically

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15

# The advantage determines if changing the policy is beneficial
 
# Consider switching from π to π' in state s:
ΔV(s) = V^{π'}(s) - V^π(s)
 
# Performance difference lemma (Kakade & Langford, 2002):
ΔV(s₀) = E_{τ~π'}[Σ_t γ^t A^π(s_t, a_t)]
 
# This says: The improvement in expected return from π to π' 
# equals the expected sum of advantages under π'!
 
# Implications:
# - If A^π(s, π'(s)) > 0 everywhere, π' is better than π
# - Policy gradient ascent exploits this by increasing probability
#   of actions with positive advantage

1
2
3
4
5
6
7
8
9
10
11
12
13
14

# Monte Carlo advantage estimate:
Â^MC_t = G_t - V(s_t)
 
where G_t = Σ_{k=0}^{T-t} γ^k r_{t+k}  (actual return)
 
# Properties:
# ✓ Unbiased (uses actual returns)
# ✗ High variance (depends on all future stochasticity)
# ✗ Requires complete episodes
 
# When to use:
# - Environments with short episodes
# - When bias is a major concern
# - As a sanity check baseline

1
2
3
4
5
6
7
8
9
10
11
12

# One-step TD advantage (also called TD error):
Â^{TD(0)}_t = r_t + γV(s_{t+1}) - V(s_t) = δ_t
 
# Properties:
# ✗ Biased (depends on V accuracy)
# ✓ Low variance (only one random step)
# ✓ Works online (no full episode needed)
 
# When to use:
# - Continuous (non-episodic) tasks
# - When V is very accurate
# - When variance dominates learning

1
2
3
4
5
6
7
8
9
10
11
12
13

# n-step advantage:
Â^{(n)}_t = Σ_{k=0}^{n-1} γ^k r_{t+k} + γ^n V(s_{t+n}) - V(s_t)
 
# Special cases:
# n = 1: TD(0), one-step bootstrap
# n = ∞: Monte Carlo, no bootstrap
 
# Properties:
# Interpolates between TD(0) and MC
# Larger n: less bias, more variance
# Smaller n: more bias, less variance
 
# Typical choices: n = 5 to 20

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18

# GAE: exponentially-weighted average of n-step advantages
 
Â^{GAE(γ,λ)}_t = Σ_{k=0}^{∞} (γλ)^k δ_{t+k}
 
where δ_t = r_t + γV(s_{t+1}) - V(s_t)
 
# Alternative form (showing it's a weighted average of n-step):
Â^{GAE}_t = (1-λ) Σ_{n=1}^{∞} λ^{n-1} Â^{(n)}_t
 
# Properties:
# λ = 0: Â = δ_t (TD(0), high bias, low variance)
# λ = 1: Â = G_t - V(s_t) (MC, low bias, high variance)
# λ ∈ (0,1): Smooth interpolation
 
# Typical choices:
# λ = 0.95 (good default)
# λ = 0.97 (less bias, more variance)
# λ = 0.90 (more bias, less variance)

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16

# Bootstrapping introduces bias via value function errors
 
# True advantage:
A^π(s, a) = E[Σ_{k=0}^∞ γ^k r_{t+k} | s_t=s, a_t=a] - V^π(s)
 
# TD(0) estimate:
Â_t = r_t + γ V̂(s_{t+1}) - V̂(s_t)
 
# Error in estimate:
Â_t - A^π(s_t, a_t) = γ(V̂(s_{t+1}) - V^π(s_{t+1})) - (V̂(s_t) - V^π(s_t))
                     + [r_t + γV^π(s_{t+1}) - E[...]]  # Variance term
 
# The first line is bias (depends on value function errors)
# The second line is variance (depends on stochasticity)
 
# Key insight: Bias compounds over time when bootstrapping!

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17

# Variance comes from stochasticity in:
# 1. Policy: Which actions are sampled
# 2. Environment: Transition dynamics
# 3. Rewards: Stochastic reward signals
 
# Monte Carlo variance (no bootstrap):
Var[G_t] = Var[Σ_{k=0}^{T-t} γ^k r_{t+k}]
         ≈ Σ_{k=0}^{T-t} γ^{2k} Var[r_{t+k}]  # (if rewards independent)
 
# This grows with horizon: O(T) variance contributions!
 
# TD(0) variance (full bootstrap):
Var[r_t + γV(s_{t+1})] = Var[r_t] + γ² Var[V(s_{t+1})]
 
# Only ONE step of variance! Much smaller.
 
# But: TD(0) assumes V is accurate, which adds bias.

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

import numpy as np
import matplotlib.pyplot as plt
 
def analyze_bias_variance(
    env,
    policy,
    value_fn,
    true_values,  # From many MC rollouts
    num_episodes=1000
):
    """
    Empirically measure bias and variance of different advantage estimators.
    """
    lambda_values = [0.0, 0.5, 0.9, 0.95, 0.99, 1.0]
    
    results = {lam: {'estimates': [], 'true': []} for lam in lambda_values}
    
    for _ in range(num_episodes):
        # Collect trajectory
        states, actions, rewards = collect_episode(env, policy)
        
        # True advantages (from many MC samples)
        true_advantages = [compute_true_advantage(s, a, true_values) 
                          for s, a in zip(states, actions)]
        
        for lam in lambda_values:
            # Compute GAE estimates
            estimates = compute_gae_advantages(
                rewards, 
                [value_fn(s) for s in states],
                gamma=0.99,
                lam=lam
            )
            
            results[lam]['estimates'].extend(estimates)
            results[lam]['true'].extend(true_advantages)
    
    # Compute bias and variance
    analysis = {}
    for lam in lambda_values:
        estimates = np.array(results[lam]['estimates'])
        true = np.array(results[lam]['true'])
        
        bias = np.mean(estimates - true)
        variance = np.var(estimates)
        mse = np.mean((estimates - true) ** 2)
        
        analysis[lam] = {
            'bias': bias,
            'variance': variance,
            'mse': mse,  # MSE = Bias² + Variance
            'bias_squared': bias ** 2
        }
    
    # Plot
    fig, axes = plt.subplots(1, 3, figsize=(15, 4))
    
    lambdas = list(analysis.keys())
    biases = [analysis[l]['bias_squared'] for l in lambdas]
    variances = [analysis[l]['variance'] for l in lambdas]
    mses = [analysis[l]['mse'] for l in lambdas]
    
    axes[0].plot(lambdas, biases, 'b-o', label='Bias²')
    axes[0].set_xlabel('λ')
    axes[0].set_ylabel('Bias²')
    axes[0].set_title('Bias² vs λ')
    
    axes[1].plot(lambdas, variances, 'r-o', label='Variance')
    axes[1].set_xlabel('λ')
    axes[1].set_ylabel('Variance')
    axes[1].set_title('Variance vs λ')
    
    axes[2].plot(lambdas, mses, 'g-o', label='MSE')
    axes[2].plot(lambdas, biases, 'b--', alpha=0.5, label='Bias²')
    axes[2].plot(lambdas, variances, 'r--', alpha=0.5, label='Variance')
    axes[2].set_xlabel('λ')
    axes[2].set_ylabel('Error')
    axes[2].set_title('MSE = Bias² + Variance')
    axes[2].legend()
    
    plt.tight_layout()
    plt.show()
    
    return analysis

1
2
3
4
5
6
7
8
9
10
11
12
13
14

# The true objective (hard to optimize directly):
J(θ) = E_{τ~π_θ}[R(τ)]
 
# The surrogate objective (easier to optimize):
L^{CPI}(θ) = E_{s~ρ^{π_old}, a~π_old}[π_θ(a|s) / π_old(a|s) · A^{π_old}(s,a)]
 
# This is the "conservative policy iteration" objective.
# Key properties:
# 1. L(θ_old) = E[A^{π_old}] = 0 (baseline)
# 2. ∇_θ L(θ)|_{θ=θ_old} = ∇_θ J(θ)|_{θ=θ_old} (same gradient!)
# 3. Easier to estimate with samples from π_old
 
# The importance weight π_θ(a|s)/π_old(a|s) corrects for
# the distribution mismatch between π_θ and π_old

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17

# The surrogate objective tells us exactly what to do:
 
L = Σ_{s,a} (π_θ(a|s) / π_old(a|s)) · A(s, a)
 
# Case 1: A(s, a) > 0 (action is good)
# To increase L, increase π_θ(a|s)
# → Make good actions MORE likely
 
# Case 2: A(s, a) < 0 (action is bad)
# To increase L, decrease π_θ(a|s)  
# → Make bad actions LESS likely
 
# Case 3: A(s, a) = 0 (action is average)
# No contribution to L
# → No need to change probability
 
# The magnitude |A(s, a)| determines how strongly to update!

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19

# TRPO: Constrained optimization
maximize L(θ)
subject to KL(π_old || π_θ) ≤ δ
 
# PPO: Clipped objective
L^{CLIP}(θ) = E[min(r_t · Â_t, clip(r_t, 1-ε, 1+ε) · Â_t)]
 
where r_t = π_θ(a|s) / π_old(a|s)
 
# For positive advantages (good actions):
# - If r > 1+ε (probability increased too much), clip
# - Prevents over-exploitation of good actions
 
# For negative advantages (bad actions):
# - If r < 1-ε (probability decreased too much), clip
# - Prevents over-penalization of bad actions
 
# This keeps policy changes conservative while still
# following the direction indicated by the advantage!

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

import torch
import numpy as np
from typing import Tuple
 
 
def compute_gae_and_returns(
    rewards: np.ndarray,            # Shape: (T,) or (T, num_envs)
    values: np.ndarray,             # Shape: (T,) or (T, num_envs)
    dones: np.ndarray,              # Shape: (T,) or (T, num_envs)
    last_value: float,              # Bootstrap value for last state
    gamma: float = 0.99,
    gae_lambda: float = 0.95
) -> Tuple[np.ndarray, np.ndarray]:
    """
    Compute Generalized Advantage Estimation and returns.
    
    Handles multiple environments and episode boundaries correctly.
    
    Args:
        rewards: Array of rewards
        values: Array of value estimates V(s_t)
        dones: Array indicating episode termination
        last_value: V(s_T) for bootstrapping (0 if terminal)
        gamma: Discount factor
        gae_lambda: GAE parameter
    
    Returns:
        advantages: GAE advantage estimates
        returns: Target returns for value training (advantages + values)
    """
    T = len(rewards)
    advantages = np.zeros_like(rewards)
    
    # Append last value for bootstrapping
    values_ext = np.append(values, last_value)
    
    # Compute backwards
    gae = 0
    for t in reversed(range(T)):
        # If episode ended, next value is 0
        if dones[t]:
            next_value = 0.0
            gae = 0  # Reset GAE at episode boundary
        else:
            next_value = values_ext[t + 1]
        
        # TD error: δ_t = r_t + γV(s_{t+1}) - V(s_t)
        delta = rewards[t] + gamma * next_value - values[t]
        
        # GAE: Â_t = δ_t + (γλ) Â_{t+1}
        gae = delta + gamma * gae_lambda * gae
        advantages[t] = gae
    
    # Returns for value function training
    returns = advantages + values
    
    return advantages, returns
 
 
class AdvantageBuffer:
    """
    Buffer that stores rollout data and computes advantages.
    
    Designed for batch policy gradient updates.
    """
    
    def __init__(
        self,
        buffer_size: int,
        state_dim: int,
        gamma: float = 0.99,
        gae_lambda: float = 0.95
    ):
        self.buffer_size = buffer_size
        self.gamma = gamma
        self.gae_lambda = gae_lambda
        
        # Pre-allocate storage
        self.states = np.zeros((buffer_size, state_dim), dtype=np.float32)
        self.actions = np.zeros(buffer_size, dtype=np.int64)
        self.rewards = np.zeros(buffer_size, dtype=np.float32)
        self.values = np.zeros(buffer_size, dtype=np.float32)
        self.log_probs = np.zeros(buffer_size, dtype=np.float32)
        self.dones = np.zeros(buffer_size, dtype=np.bool_)
        
        # Computed after rollout
        self.advantages = np.zeros(buffer_size, dtype=np.float32)
        self.returns = np.zeros(buffer_size, dtype=np.float32)
        
        self.ptr = 0
    
    def store(
        self,
        state: np.ndarray,
        action: int,
        reward: float,
        value: float,
        log_prob: float,
        done: bool
    ):
        """Store a single transition."""
        idx = self.ptr
        self.states[idx] = state
        self.actions[idx] = action
        self.rewards[idx] = reward
        self.values[idx] = value
        self.log_probs[idx] = log_prob
        self.dones[idx] = done
        self.ptr += 1
    
    def finish_rollout(self, last_value: float):
        """
        Compute advantages after rollout is complete.
        
        Call this before getting data for training.
        """
        self.advantages, self.returns = compute_gae_and_returns(
            self.rewards[:self.ptr],
            self.values[:self.ptr],
            self.dones[:self.ptr],
            last_value,
            self.gamma,
            self.gae_lambda
        )
        
        # Normalize advantages
        self.advantages = (
            (self.advantages - self.advantages.mean()) / 
            (self.advantages.std() + 1e-8)
        )
    
    def get_data(self) -> dict:
        """Get all data as tensors for training."""
        return {
            'states': torch.FloatTensor(self.states[:self.ptr]),
            'actions': torch.LongTensor(self.actions[:self.ptr]),
            'log_probs': torch.FloatTensor(self.log_probs[:self.ptr]),
            'advantages': torch.FloatTensor(self.advantages[:self.ptr]),
            'returns': torch.FloatTensor(self.returns[:self.ptr])
        }
    
    def reset(self):
        """Clear buffer for next rollout."""
        self.ptr = 0

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

# Standard normalization (recommended)
def normalize_advantages(advantages: torch.Tensor) -> torch.Tensor:
    """Zero-mean, unit-variance normalization."""
    return (advantages - advantages.mean()) / (advantages.std() + 1e-8)
 
 
# Per-batch normalization (used in most implementations)
# Applied to each minibatch during training
def normalize_batch(batch_advantages: torch.Tensor) -> torch.Tensor:
    mean = batch_advantages.mean()
    std = batch_advantages.std()
    return (batch_advantages - mean) / (std + 1e-8)
 
 
# Alternative: Normalize over entire rollout, not per-batch
# Can be more stable but less adaptive
class RunningNormalizer:
    """Track running mean and std for consistent normalization."""
    
    def __init__(self, momentum: float = 0.99):
        self.mean = 0.0
        self.var = 1.0
        self.momentum = momentum
        self.count = 0
    
    def update(self, batch: np.ndarray):
        batch_mean = batch.mean()
        batch_var = batch.var()
        batch_count = len(batch)
        
        self.mean = self.momentum * self.mean + (1 - self.momentum) * batch_mean
        self.var = self.momentum * self.var + (1 - self.momentum) * batch_var
        self.count += batch_count
    
    def normalize(self, x: np.ndarray) -> np.ndarray:
        return (x - self.mean) / (np.sqrt(self.var) + 1e-8)
 
 
# When NOT to normalize (rare cases):
# 1. When advantage magnitudes carry important signal
# 2. When batch size is very small (unstable statistics)
# 3. When debugging/understanding raw advantage 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

class DuelingNetwork(nn.Module):
    """
    Learn V(s) and A(s,a) separately, combine for Q(s,a).
    
    From: Dueling Network Architectures (Wang et al., 2016)
    
    Advantages of this factorization:
    1. V(s) can be learned from any action (data efficient)
    2. A(s,a) focuses on relative action quality
    3. Often faster learning, especially when many actions similar
    """
    
    def __init__(self, state_dim, action_dim, hidden_dim=256):
        super().__init__()
        
        # Shared feature extraction
        self.features = nn.Sequential(
            nn.Linear(state_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU()
        )
        
        # Value stream: V(s)
        self.value_stream = nn.Sequential(
            nn.Linear(hidden_dim, hidden_dim // 2),
            nn.ReLU(),
            nn.Linear(hidden_dim // 2, 1)
        )
        
        # Advantage stream: A(s, a)
        self.advantage_stream = nn.Sequential(
            nn.Linear(hidden_dim, hidden_dim // 2),
            nn.ReLU(),
            nn.Linear(hidden_dim // 2, action_dim)
        )
    
    def forward(self, state):
        features = self.features(state)
        
        value = self.value_stream(features)  # [batch, 1]
        advantages = self.advantage_stream(features)  # [batch, actions]
        
        # Combine: Q(s,a) = V(s) + (A(s,a) - mean(A))
        # Subtracting mean ensures advantages are centered
        # and V(s) = E_a[Q(s,a)] holds
        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

# Retrace(λ): Corrects for off-policy data
 
# Problem: When using data from old policy π_old, 
# advantages estimated for current π are biased
 
# Retrace correction:
Â^{Retrace}_t = δ_t + γλ · c_{t+1} · Â^{Retrace}_{t+1}
 
where c_t = min(1, π(a_t|s_t) / π_old(a_t|s_t))
 
# The correction factor c_t:
# - Truncates importance weights to reduce variance
# - Allows safe off-policy learning
# - Used in algorithms like ACER, Reactor
 
# V-trace (from IMPALA):
# Similar idea but applied to value function learning
# Enables large-scale distributed RL with stale data

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16

# Intrinsic motivation augments advantages with curiosity bonuses
 
# Standard advantage:
A(s, a) = Q(s, a) - V(s)
 
# With intrinsic motivation:
A^{total}(s, a) = A^{extrinsic}(s, a) + β · A^{intrinsic}(s, a)
 
# Intrinsic reward examples:
# - Prediction error: How surprising was the transition?
# - Novelty: How rarely has this state been visited?
# - Information gain: How much did we learn?
 
# Random Network Distillation (RND):
r^{intrinsic} = ||f_target(s') - f_predictor(s')||²
# High when s' is novel (predictor hasn't seen it before)