N-Step Tree Backup One Node

Problem

Compute a two-step tree-backup target.

The target is r0 + gamma * sum_a pi1[a] * (r1[a] + gamma * sum_b pi2[a][b] * q2[a][b]), where pi2 and q2 are indexed first by the action at the next state and then by the following action.

Hint

Every branch of the tree is weighted by the target policy, never by the behavior policy.

Test cases

case 1

in: [0,0.9,[0.5,0.5],[1,2],[[0.5,0.5],[1,0]],[[2,4],[3,0]]]

expected: 3.7800000000000002

case 2

in: [1,0.5,[1,0],[2,0],[[0.25,0.75],[0.5,0.5]],[[4,8],[1,1]]]

expected: 3.75

case 3

in: [0.5,1,[0.5,0.5],[1,1],[[1,0],[0,1]],[[2,1],[1,2]]]

expected: 3.5

Starter codePython

def tree_backup_two_step(r0, gamma, pi1, r1, pi2, q2):
    # Your code here
    pass

3 test cases · run your code against them in the editor