ICM Inverse Dynamics Loss

Problem

Compute the ICM inverse dynamics loss for a single transition: the negative log-softmax of the true action under the predicted logits. Return -log(softmax(logits)[action]) computed with the max-subtraction trick.

Hint

The inverse model learns which action caused the observed feature change.

Test cases

case 1

in: [[0,0],0]

expected: 0.6931471805599453

case 2

in: [[2,0],0]

expected: 0.1269280110429726

case 3

in: [[1,2,3],1]

expected: 1.4076059644443804

Starter codePython

import math
def icm_inverse_loss(logits, action):
    # Your code here
    pass

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