Decision-Transformer Return Conditioning

Problem

Build return-to-go conditioning tokens for a decision transformer.

At each step the token is target_return minus the rewards collected so far, so return [target_return, target_return - r_0, target_return - r_0 - r_1, ...].

Hint

Conditioning on return-to-go lets the transformer ask for higher-return behavior at test time.

Test cases

case 1

in: [100,[10,20,30]]

expected: [100,90,70]

case 2

in: [5,[1,1]]

expected: [5,4]

case 3

in: [10,[]]

expected: []

Starter codePython

def return_to_go_conditioning(target_return, rewards):
    # Your code here
    pass

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