Bellman Optimality Backup on Action Values

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.