Bellman Optimality Backup on Action Values

mediumPROMarkov Decision Processes and Bellman EquationsReinforcement LearningRL 1: Foundations

Problem Statement

Apply one sweep of the Bellman optimality operator to a table of action values.

Background

The Bellman expectation equation averages over whatever the policy does. The Bellman optimality equation replaces that average with a maximum, which is what makes it non-linear and what makes it describe the best achievable behaviour rather than some particular behaviour:

$q_*(s,a) = r(s,a) + \gamma \sum_{s'} p(s' \mid s, a)\, \max_{a'} q_*(s', a')$

Read the structure carefully. The action a in state s is given — you are committed to it, so no max there. The max lives at the successor state, encoding "after this action I will act optimally from then on". Turning that equation into an assignment gives the optimality operator, and repeatedly applying it is value iteration.

This operator is a $\gamma$ -contraction in the max norm, so for gamma < 1 it has a unique fixed point and iterating from any starting table converges to $q_*$ . And because the max is attained by some action in every state, a deterministic optimal policy always exists — you never need to randomise to be optimal in a finite MDP.

Your Task

Implement:

def q_optimality_backup(P, R, q, gamma):
    ...

P[s][a][s2] — transition probabilities.
R[s][a] — expected immediate reward.
q[s][a] — the current action-value table.
gamma — discount factor.

Return the new table as a list of n_states lists of n_actions floats. Use the old q for every entry of the sweep (a synchronous update); do not read back values you wrote during this same sweep.

Input / Output Format

Nested lists of floats in, a nested list of floats out, rounded to 4 decimals by the grader.

Sample

P = [[[0.0, 1.0], [1.0, 0.0]], [[1.0, 0.0], [0.0, 1.0]]]
R = [[0.0, 1.0], [2.0, 0.0]]
q = [[1.0, 3.0], [4.0, -1.0]]
print([[round(x, 4) for x in row] for row in q_optimality_backup(P, R, q, 0.5)])

Output:

[[2.0, 2.5], [3.5, 2.0]]

The successor-state maxima are max(q[0]) = 3.0 and max(q[1]) = 4.0. Then q'(0,0) = 0.0 + 0.5*4.0 = 2.0, q'(0,1) = 1.0 + 0.5*3.0 = 2.5, and so on.

Example

Input

q_optimality_backup([[[0.0, 1.0], [1.0, 0.0]], [[1.0, 0.0], [0.0, 1.0]]], [[0.0, 1.0], [2.0, 0.0]], [[1.0, 3.0], [4.0, -1.0]], 0.5)

Output

[[2.0, 2.5], [3.5, 2.0]]

Reasoning

First take the max over actions at each successor state: max(q[0])=3.0, max(q[1])=4.0. Action 0 in state 0 lands in state 1 with certainty, so q'(0,0)=0.0+0.54.0=2.0. Action 1 in state 0 lands in state 0, so q'(0,1)=1.0+0.53.0=2.5.

Constraints

1 <= n_states <= 100, 1 <= n_actions <= 20
Each P[s][a] sums to 1.
0.0 <= gamma <= 1.0
The update must be synchronous: compute every new entry from the old q table.
Do not round inside the function.

Learn this in

Reinforcement LearningML Study Plan Markov Decision ProcessesReinforcement Learning Dynamic ProgrammingReinforcement Learning

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

Concept primerPREVIEW

1. Background Knowledge

The Bellman optimality equation is the cornerstone of value-based reinforcement learning. It defines the optimal action-value function, denoted as $q_*(s, a)$ , which represents the maximum expected return achievable from state $s$ by taking action $a$ and thereafter behaving optimally. Unlike the Bellman expectation equation, which averages over a specific policy, the optimality equation uses a maximum operator over all possible actions in the successor state. This non-linear operator captures the essence of optimal decision-making: after taking the current action, the agent will choose the best possible action in the next state.

The mathematical formulation is given by:

q_*(s,a) = r(s,a) + \gamma \sum_{s'} p(s' \mid s, a)\, \max_{a'} q_*(s', a')

Here, $r(s,a)$ is the immediate reward, $\gamma$ is the discount factor, and $p(s' \mid s, a)$ is the probability of transitioning to state $s'$ . The term $\max_{a'} q_*(s', a')$ is crucial: it signifies that once the transition occurs, the agent acts greedily with respect to the optimal value function. This equation forms the basis for value iteration, an algorithm that converges to $q_*$ by repeatedly applying this backup operation.

In practical implementations, we often work with a table of action values, $q$ , rather than the true $q_*$ . The Bellman optimality backup is the operation that updates this table. A single sweep of this backup computes a new estimate for every state-action pair using the current estimates of the successor states. This process is a contraction mapping, meaning that repeated application brings the value estimates closer to the true optimal values, guaranteeing convergence for $\gamma < 1$ .

2. Algorithm Approach

The core approach is to implement a synchronous update of the action-value table. This means we must compute all new values based on the old values from the previous iteration, not the newly computed ones from the current iteration. This distinction is vital for the theoretical convergence properties of value iteration.

The algorithm follows these high-level steps:

Initialize a new table q_new with the same dimensions as the input q.
Iterate through every state $s$ and every action $a$ .
Compute the immediate reward $R[s][a]$ .
Calculate the expected maximum value of the successor states. This involves:

Iterating through all possible next states $s'$ .
Finding the maximum action value $\max_{a'} q[s'][a']$ for each successor state $s'$ .
Weighting these maxima by their transition probabilities $P[s][a][s']$ .
Summing these weighted maxima.

Apply the Bellman equation: $q_{new}[s][a] = R[s][a] + \gamma \times (\text{expected max successor value})$ .
Return the q_new table.

This approach ensures that the update is consistent with the definition of the Bellman optimality operator, treating the current state-action pair as fixed and optimizing only for the future.

3. Step-by-Step Strategy

Determine Dimensions: Identify the number of states ( $n\_states$ ) and actions ( $n\_actions$ ) from the input q or P. Initialize q_new as a list of lists with the same structure, filled with zeros or copies of q (though zeros are safer to avoid accidental mixing).

PRO

Continue the full explanation

You're reading the free preview. Unlock the complete walkthrough, the code editor, test runner and reference solution with Premium.

PRO

Editor locked

The code editor is locked for Pro problems. It is only available for free problems. Please upgrade to gain access to the code editor for all problems.

solution.py

Python 3.13

Test results