In-Context RL Pick

Problem

Pick the action from the demonstration whose state is nearest to the query state (Euclidean distance, ties toward the earlier demonstration). No demonstrations returns -1.

Hint

Nearest-neighbor retrieval is the simplest form of in-context decision making.

Test cases

case 1

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

expected: 1

case 2

in: [[4,4],[[0,0],[5,5]],[0,1]]

expected: 1

case 3

in: [[1],[],[]]

expected: -1

Starter codePython

def in_context_rl_pick(query_state, demo_states, demo_actions):
    # Your code here
    pass

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