📘

SARSA, Q-Learning and Expected SARSA Targets

MediumReinforcement Learning 📚 RL 1: Foundations

Problem Statement

Implement the three action-value control updates side by side. They differ only in how the next state is valued.

Background

Every one of these has the same skeleton — move Q(s,a) a fraction alpha toward a target — and the entire on-policy/off-policy distinction lives in one term.

SARSA uses the action the behaviour policy actually took next:

$Q(S_t,A_t) \leftarrow Q(S_t,A_t) + \alpha\big[R_{t+1} + \gamma Q(S_{t+1},A_{t+1}) - Q(S_t,A_t)\big]$

It is on-policy: it learns the value of the policy you are running, exploration mistakes included. That is why SARSA learns the "safe" path near a cliff — it knows an epsilon-greedy agent will occasionally fall off.

Q-learning ignores what happened next and uses the greedy value:

$Q(S_t,A_t) \leftarrow Q(S_t,A_t) + \alpha\big[R_{t+1} + \gamma \max_{a'} Q(S_{t+1},a') - Q(S_t,A_t)\big]$

It is off-policy: it learns the optimal policy's values while behaving with exploration. The price is maximisation bias — a max over noisy estimates is systematically optimistic.

Expected SARSA replaces the sample or the max with the exact expectation under the target policy:

$Q(S_t,A_t) \leftarrow Q(S_t,A_t) + \alpha\Big[R_{t+1} + \gamma \sum_{a'} \pi(a' \mid S_{t+1}) Q(S_{t+1},a') - Q(S_t,A_t)\Big]$

Removing the sampling of $A_{t+1}$ removes its variance for the cost of one extra sum, which is usually free. With an epsilon-greedy $\pi$ and epsilon = 0 it is Q-learning; with larger epsilon it drifts toward the on-policy answer.

Your Task

Implement all three. Each returns the new scalar value of Q[s][a]; do not mutate Q.

def sarsa_update(Q, s, a, r, s2, a2, alpha, gamma):
    ...

def q_learning_update(Q, s, a, r, s2, alpha, gamma):
    ...

def expected_sarsa_update(Q, s, a, r, s2, alpha, gamma, epsilon):
    ...

expected_sarsa_update weights the successor action values by an **epsilon-greedy distribution over **Q[s2]****, splitting the greedy probability 1 - epsilon evenly across all tied maxima.

Input / Output Format

Q is a nested list of floats. Each function returns a single float, rounded to 4 decimals by the grader.

Sample

Q = [[0.0, 0.0], [2.0, 6.0]]
print(round(sarsa_update(Q, 0, 0, 1.0, 1, 0, 0.5, 0.9), 4))
print(round(q_learning_update(Q, 0, 0, 1.0, 1, 0.5, 0.9), 4))
print(round(expected_sarsa_update(Q, 0, 0, 1.0, 1, 0.5, 0.9, 0.2), 4))

Output:

1.4
3.2
3.02

All three start from Q[0][0] = 0.0 and move half way (alpha = 0.5) toward their target. SARSA bootstraps off the action actually taken, Q[1][0] = 2.0, for a target of 1.0 + 0.9*2.0 = 2.8. Q-learning bootstraps off max(Q[1]) = 6.0, for a target of 6.4. Expected SARSA bootstraps off the epsilon-greedy expectation 0.12.0 + 0.96.0 = 5.6, for a target of 1.0 + 0.9*5.6 = 6.04.

Example:

Input:

Q = [[0.0, 0.0], [2.0, 6.0]]; sarsa_update(Q, 0, 0, 1.0, 1, 0, 0.5, 0.9), q_learning_update(Q, 0, 0, 1.0, 1, 0.5, 0.9), expected_sarsa_update(Q, 0, 0, 1.0, 1, 0.5, 0.9, 0.2)

Output:

1.4
3.2
3.02

Reasoning:

All three move Q(0,0)=0 half way toward their target because alpha=0.5. SARSA uses the action actually taken next, a2=0, so its target is 1.0 + 0.92.0 = 2.8 and the result is 1.4. Q-learning uses max(Q[1]) = 6.0 for a target of 6.4, giving 3.2. Expected SARSA weights the successor values by the epsilon-greedy distribution: with epsilon=0.2 over 2 actions the greedy action gets 0.8 + 0.1 = 0.9 and the other gets 0.1, so the expectation is 0.12.0 + 0.96.0 = 5.6, the target is 1.0 + 0.95.6 = 6.04, and the result is 3.02.

Constraints:

1 <= n_states <= 500, 1 <= n_actions <= 50
0.0 <= alpha <= 1.0, 0.0 <= gamma <= 1.0, 0.0 <= epsilon <= 1.0
Q must not be modified — return the new scalar only.
Ties for the max in the epsilon-greedy distribution split 1 - epsilon evenly.
Do not round inside the functions.

More from RL 1: Foundations

GAE Over a Full TrajectoryHard PPO Clipped Objective (Single Sample)Medium TD Residuals for GAEMedium Return Standardization BaselineMedium REINFORCE Loss TermMedium

1. Background Knowledge

Temporal-Difference (TD) Learning is a core reinforcement learning technique that combines Monte Carlo ideas (learning from experience) with dynamic programming (bootstrapping). The fundamental update rule adjusts an estimate $Q(s, a)$ toward a target value using a learning rate $\alpha$ . The general form is $Q_{new} = Q_{old} + \alpha(\text{Target} - Q_{old})$ . The critical distinction between algorithms lies solely in how this Target is constructed.

SARSA is an on-policy algorithm. It learns the value of the policy currently being executed. Its target uses the actual next action $A_{t+1}$ taken by the agent: $R_{t+1} + \gamma Q(S_{t+1}, A_{t+1})$ . Because it accounts for exploration noise (e.g., $\epsilon$ -greedy mistakes), SARSA tends to learn safer policies that avoid high-risk areas if the agent might accidentally fall into them.

Q-Learning is an off-policy algorithm. It learns the optimal action-value function regardless of the agent's behavior. Its target uses the maximum possible future value: $R_{t+1} + \gamma \max_{a'} Q(S_{t+1}, a')$ . This assumes the agent will act optimally in the next step, ignoring the exploration strategy used to reach $S_{t+1}$ . This can lead to maximization bias, where noisy estimates are systematically overestimated.

Expected SARSA bridges the gap by using the expected value of the next state under the current policy $\pi$ . Instead of sampling a single next action or taking the max, it computes $\sum_{a'} \pi(a' \mid S_{t+1}) Q(S_{t+1}, a')$ . This reduces variance compared to SARSA while remaining consistent with the on-policy behavior. For an $\epsilon$ -greedy policy, the probability of taking the greedy action(s) is $1-\epsilon$ (split evenly among ties), and the probability of any other action is $\epsilon / |A|$ .

2. Algorithm Approach

The approach for all three functions is identical in structure: calculate the TD Target and then apply the TD Update formula. You must not modify the input Q table; instead, compute the new value for Q[s][a] and return it.

Retrieve Current Value: Get $Q_{old} = Q[s][a]$ .
Compute Target:

SARSA: Target $= r + \gamma \cdot Q[s2][a2]$
Q-Learning: Target $= r + \gamma \cdot \max(Q[s2])$
Expected SARSA: Target $= r + \gamma \cdot \text{ExpectedValue}(Q[s2], \epsilon)$

Apply Update: $Q_{new} = Q_{old} + \alpha \cdot (\text{Target} - Q_{old})$
Return: Return $Q_{new}$ .

The complexity lies entirely in correctly implementing the Expected Value calculation for Expected SARSA, specifically handling the $\epsilon$ -greedy probability distribution over tied maximum actions.

3. Step-by-Step Strategy

For sarsa_update and q_learning_update:

Extract the current Q-value: q_old = Q[s][a].
For SARSA, the next state value is simply Q[s2][a2].
For Q-Learning, find the maximum value in the list Q[s2]. Use max(Q[s2]).
Calculate the target: target = r + gamma * next_state_value.
Calculate the update: q_new = q_old + alpha * (target - q_old).
Return q_new.

For expected_sarsa_update:

Extract the current Q-value: q_old = Q[s][a].
Identify the maximum Q-value in the next state s2: max_q = max(Q[s2]).
Count how many actions in Q[s2] equal max_q. Let this count be num_ties.
Calculate the probability of taking the greedy action(s): p_greedy = (1 - epsilon) / num_ties.
Calculate the probability of taking any non-greedy action: p_random = epsilon / len(Q[s2]).
Compute the expected value:

Sum of all Q-values in Q[s2] multiplied by p_random.
Add the correction for greedy actions: For each action that is a tie for max, add (p_greedy - p_random) * max_q.
Alternatively, simpler logic: expected_val = sum(Q[s2]) * (epsilon / n) + max_q * ((1 - epsilon) / num_ties). Wait, this double counts if not careful.
Correct Logic: The expected value is $\sum_{a'} \pi(a'|s2) Q(s2, a')$ .
$\pi(a'|s2) = \frac{\epsilon}{|A|}$ for all $a'$ .
For actions that are greedy (max), add extra probability: $\frac{1-\epsilon}{\text{num\_ties}}$ .
So, expected_val = sum(Q[s2]) * (epsilon / n) + max_q * ((1 - epsilon) / num_ties).
Actually, a cleaner way: Start with the uniform random expectation: base = sum(Q[s2]) / n. Then adjust for the greedy bias.
Easiest implementation: Iterate through all actions in Q[s2]. If Q[s2][a'] == max_q, its probability is (1-epsilon)/num_ties. Else, it is epsilon/n. Sum prob * value.

Calculate target: target = r + gamma * expected_val.
Calculate update: q_new = q_old + alpha * (target - q_old).
Return q_new.

4. Common Pitfalls

Mutating Input: The problem explicitly states "do not mutate Q". Ensure you only read from Q and return a new float. Do not perform Q[s][a] =....
Tied Maxima in Expected SARSA: If multiple actions in Q[s2] have the same maximum value, the probability $1-\epsilon$ must be split evenly among them. Failing to divide by num_ties is a common error.
Probability Normalization: Ensure probabilities sum to 1. For Expected SARSA, verify that the sum of probabilities for all actions in s2 equals 1.
Floating Point Precision: While the grader rounds to 4 decimals, intermediate calculations should use standard float precision. Do not round intermediate steps.
Indexing Errors: Ensure s2 is a valid index for Q and a2 is a valid index for Q[s2].
Confusing SARSA and Q-Learning Targets: Remember SARSA uses the actual next action a2, while Q-Learning uses the best possible next action max(Q[s2]).

5. Time & Space Complexity

Time Complexity:
sarsa_update: $O(1)$ as it accesses specific indices.
q_learning_update: $O(|A|)$ where $|A|$ is the number of actions, due to finding the max in Q[s2].
expected_sarsa_update: $O(|A|)$ to find the max, count ties, and sum the expected value.
Since $|A|$ is typically small in grid-world problems, this is effectively constant time per update.
Space Complexity:
$O(1)$ auxiliary space for all functions. We only store a few scalar variables (q_old, target, expected_val, etc.) and do not allocate new data structures proportional to the input size.

solution.py

Python 3.13

Test Results

0/0

Run code to see test results.