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

# Multi-Armed Bandit: k slot machines (arms)
# Each arm i has unknown reward distribution with mean μ_i
# Goal: maximize total reward over T pulls
 
class MultiArmedBandit:
    def __init__(self, k_arms):
        # True (unknown) mean rewards for each arm
        self.true_means = np.random.randn(k_arms)
    
    def pull(self, arm):
        # Returns noisy reward from arm's distribution
        return self.true_means[arm] + np.random.randn()
 
# The challenge: Find arm with highest μ_i quickly
# Trade off exploring to find best arm vs. exploiting known good arms
 
# Regret: Difference between optimal strategy and our strategy
# R(T) = T * μ* - Σ_{t=1}^{T} μ_{a_t}
#       ↑ Best possible    ↑ What we got
 
# Good algorithms achieve sublinear regret: R(T) = O(log T)
# This means exploration overhead becomes negligible as T → ∞

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

def epsilon_greedy_action(q_values, epsilon):
    """
    With probability epsilon: random action
    With probability (1-epsilon): best action according to Q
    """
    if np.random.random() < epsilon:
        return np.random.randint(len(q_values))  # Explore
    else:
        return np.argmax(q_values)                # Exploit
 
# In practice: decay epsilon over training
def get_epsilon(step, epsilon_start=1.0, epsilon_end=0.01, decay_steps=100000):
    """Linear decay from epsilon_start to epsilon_end."""
    progress = min(step / decay_steps, 1.0)
    return epsilon_start - progress * (epsilon_start - epsilon_end)
 
# DQN uses: epsilon 1.0 → 0.1 over first 1M steps, then stays at 0.1

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

def boltzmann_action(q_values, temperature):
    """
    Sample action proportional to exp(Q/τ).
    Temperature τ controls exploration:
    - High τ → uniform (exploration)
    - Low τ → greedy (exploitation)
    """
    # Compute softmax probabilities
    scaled_q = q_values / temperature
    scaled_q -= scaled_q.max()  # For numerical stability
    exp_q = np.exp(scaled_q)
    probs = exp_q / exp_q.sum()
    
    return np.random.choice(len(q_values), p=probs)
 
# Advantage over ε-greedy: 
# Actions with higher Q are more likely even during exploration
# Actions with similar Q get similar probabilities

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

def ucb_action(q_values, action_counts, total_count, c=2.0):
    """
    Select action that maximizes: Q(a) + c * sqrt(log(t) / n(a))
                                   ↑       ↑
                                 exploit  explore
    
    Actions with low counts get exploration bonus.
    c controls exploration-exploitation balance.
    """
    exploration_bonus = c * np.sqrt(np.log(total_count) / (action_counts + 1e-8))
    ucb_values = q_values + exploration_bonus
    return np.argmax(ucb_values)
 
# Why it works:
# - Uncertainty shrinks as sqrt(1/n) for each action
# - Unexplored actions have high uncertainty → high bonus
# - As action is tried more, bonus decreases → exploit
 
# Provably achieves O(log T) regret in multi-armed bandits

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

class ThompsonSampling:
    """
    Bayesian approach: Maintain posterior over reward distributions.
    Sample from posteriors and act greedily on sample.
    """
    
    def __init__(self, n_arms):
        # Beta distribution parameters for each arm
        # Beta(α, β) models probability of success
        self.alphas = np.ones(n_arms)  # Successes + 1
        self.betas = np.ones(n_arms)   # Failures + 1
    
    def select_action(self):
        # Sample from posterior for each arm
        sampled_probs = np.random.beta(self.alphas, self.betas)
        return np.argmax(sampled_probs)
    
    def update(self, arm, reward):
        # Update posterior based on observed reward
        if reward > 0.5:  # Assuming binary reward
            self.alphas[arm] += 1
        else:
            self.betas[arm] += 1
 
# Why Thompson Sampling works:
# - Naturally balances exploration (uncertain arms have wide posteriors)
# - Automatically decreases exploration as uncertainty shrinks
# - Often outperforms UCB empirically
# - Computationally simple for conjugate priors

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

# Total reward = Extrinsic (from environment) + Intrinsic (from curiosity)
 
def combined_reward(state, action, next_state, env_reward, intrinsic_module):
    """
    Combine external task reward with intrinsic exploration reward.
    """
    extrinsic = env_reward                                    # From environment
    intrinsic = intrinsic_module.compute_bonus(state, action, next_state)
    
    # β controls exploration vs. exploitation
    total = extrinsic + beta * intrinsic
    return total
 
# Common pattern:
# - Early training: β high → explore widely
# - Later training: β decays → focus on task reward
 
# Key design choice: What should intrinsic reward measure?
# Options: Novelty, Surprise, Prediction Error, Information Gain

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

class PseudoCountExploration:
    """
    Use density model ρ(s) to compute pseudo-counts.
    Pseudo-count: N̂(s) = ρ(s) / (1 - ρ(s))
    
    Intuition: If ρ(s) is low, state is novel → high bonus.
    """
    
    def __init__(self, state_dim):
        # Train a density model (e.g., PixelCNN for images)
        self.density_model = DensityModel(state_dim)
    
    def compute_bonus(self, state):
        # Get density before and after seeing this state
        rho_old = self.density_model.predict_density(state)
        self.density_model.update(state)
        rho_new = self.density_model.predict_density(state)
        
        # Compute pseudo-count
        if rho_old == 0:
            return 1.0  # Completely novel state
        
        pseudo_count = rho_old / max(1e-10, rho_new - rho_old)
        
        # Exploration bonus inversely proportional to pseudo-count
        bonus = 1.0 / np.sqrt(pseudo_count + 1e-8)
        return bonus

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

class ICM(nn.Module):
    """
    Intrinsic Curiosity Module
    Curiosity = prediction error in learned feature space
    
    Key insight: Inverse model forces features to capture 
    only aspects affected by actions, filtering out noise.
    """
    
    def __init__(self, obs_dim, action_dim, feature_dim=256):
        super().__init__()
        
        # Feature encoder: obs → features
        self.encoder = nn.Sequential(
            nn.Conv2d(3, 32, 8, stride=4),
            nn.ReLU(),
            nn.Conv2d(32, 64, 4, stride=2),
            nn.ReLU(),
            nn.Conv2d(64, 64, 3, stride=1),
            nn.ReLU(),
            nn.Flatten(),
            nn.Linear(64 * 7 * 7, feature_dim)
        )
        
        # Forward model: (φ(s), a) → predicted φ(s')
        self.forward_model = nn.Sequential(
            nn.Linear(feature_dim + action_dim, 256),
            nn.ReLU(),
            nn.Linear(256, feature_dim)
        )
        
        # Inverse model: (φ(s), φ(s')) → predicted action
        self.inverse_model = nn.Sequential(
            nn.Linear(feature_dim * 2, 256),
            nn.ReLU(),
            nn.Linear(256, action_dim)
        )
    
    def compute_intrinsic_reward(self, state, action, next_state):
        """
        Intrinsic reward = forward prediction error in feature space.
        """
        phi_s = self.encoder(state)
        phi_next_s = self.encoder(next_state)
        
        # One-hot encode action
        action_onehot = F.one_hot(action, num_classes=self.action_dim)
        
        # Predict next features from current features + action
        predicted_phi_next = self.forward_model(
            torch.cat([phi_s, action_onehot], dim=-1)
        )
        
        # Curiosity = prediction error (squared L2 norm)
        intrinsic_reward = 0.5 * ((predicted_phi_next - phi_next_s) ** 2).sum(dim=-1)
        
        return intrinsic_reward
    
    def compute_losses(self, state, action, next_state):
        """Training losses for ICM."""
        phi_s = self.encoder(state)
        phi_next_s = self.encoder(next_state)
        
        # Forward model loss
        predicted_phi_next = self.forward_model(
            torch.cat([phi_s, F.one_hot(action, self.action_dim)], dim=-1)
        )
        forward_loss = F.mse_loss(predicted_phi_next, phi_next_s.detach())
        
        # Inverse model loss
        predicted_action = self.inverse_model(torch.cat([phi_s, phi_next_s], dim=-1))
        inverse_loss = F.cross_entropy(predicted_action, action)
        
        return forward_loss, inverse_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

class RND(nn.Module):
    """
    Random Network Distillation
    Intrinsic reward = prediction error matching a fixed random network
    
    Brilliantly simple: random targets + prediction error = novelty detection.
    """
    
    def __init__(self, obs_dim, feature_dim=512):
        super().__init__()
        
        # Target network: FIXED random weights (never trained!)
        self.target = nn.Sequential(
            nn.Linear(obs_dim, 512),
            nn.ReLU(),
            nn.Linear(512, 512),
            nn.ReLU(),
            nn.Linear(512, feature_dim)
        )
        # Freeze target network
        for param in self.target.parameters():
            param.requires_grad = False
        
        # Predictor network: trained to match target
        self.predictor = nn.Sequential(
            nn.Linear(obs_dim, 512),
            nn.ReLU(),
            nn.Linear(512, 512),
            nn.ReLU(),
            nn.Linear(512, feature_dim)
        )
    
    def compute_intrinsic_reward(self, obs):
        """
        Novel states: predictor hasn't learned target's output → high error.
        Familiar states: predictor has learned target's output → low error.
        """
        with torch.no_grad():
            target_features = self.target(obs)
        predicted_features = self.predictor(obs)
        
        # MSE as intrinsic reward
        intrinsic_reward = ((predicted_features - target_features) ** 2).sum(dim=-1)
        return intrinsic_reward
    
    def update(self, obs):
        """Train predictor to match target on visited states."""
        target_features = self.target(obs).detach()
        predicted_features = self.predictor(obs)
        loss = F.mse_loss(predicted_features, target_features)
        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

class GoExplore:
    """
    Go-Explore: Remember promising states, return to them, explore further.
    
    Two phases:
    1. Exploration Phase: Find goals and trajectories in deterministic setting
    2. Robustification Phase: Learn policy that works under stochasticity
    """
    
    def __init__(self, env, state_encoder):
        self.env = env
        self.state_encoder = state_encoder  # Maps states to discrete cells
        self.archive = {}  # cell → (state, trajectory, score)
    
    def exploration_phase(self, num_iterations):
        """Phase 1: Explore using archive-based search."""
        for _ in range(num_iterations):
            # 1. SELECT: Choose cell to explore from (prefer under-explored)
            cell = self.select_cell()
            state, trajectory, _ = self.archive[cell]
            
            # 2. GO: Return to the selected state
            self.env.reset()
            for action in trajectory:
                self.env.step(action)  # Deterministic playback
            
            # 3. EXPLORE: Take random actions from this state
            for step in range(explore_steps):
                action = self.env.action_space.sample()
                next_state, reward, done, _ = self.env.step(action)
                trajectory.append(action)
                
                # Encode state to cell
                new_cell = self.state_encoder(next_state)
                
                # Update archive if this is a new or better path to cell
                if new_cell not in self.archive or len(trajectory) < len(self.archive[new_cell][1]):
                    self.archive[new_cell] = (next_state, trajectory.copy(), reward)
                
                if done:
                    break
    
    def robustification_phase(self, trajectories):
        """
        Phase 2: Learn robust policy from demonstrated trajectories.
        Uses imitation learning + RL to handle stochasticity.
        """
        # Train goal-conditioned policy to reach discovered states
        policy = train_goal_conditioned_policy(trajectories)
        return policy

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

# DIAYN (Diversity is All You Need)
# Learn diverse skills without extrinsic reward
# Skills that are distinguishable → diverse exploration
 
# Objective: Maximize mutual information I(s; z)
# Intuition: Each skill z should visit distinctive states
 
class DIAYN:
    def __init__(self, policy, discriminator, n_skills):
        self.policy = policy           # π(a|s, z)
        self.discriminator = discriminator  # q(z|s)
        self.n_skills = n_skills
    
    def compute_intrinsic_reward(self, state, skill):
        """
        Reward for visiting states that identify the skill.
        High reward = discriminator correctly identifies skill from state.
        """
        # Probability discriminator assigns to the correct skill
        skill_probs = self.discriminator(state)
        log_q_z_given_s = torch.log(skill_probs[skill] + 1e-8)
        
        # Subtract log prior to encourage diversity
        log_p_z = np.log(1.0 / self.n_skills)
        
        intrinsic_reward = log_q_z_given_s - log_p_z
        return intrinsic_reward