World Model Rollout Tiny

Problem

Roll a deterministic linear world model forward for a fixed number of steps.

next = W * state + b at each step, with each prediction fed back as the next input. Return the trajectory including the initial state, so its length is steps + 1.

Hint

Feed each prediction back as the next input.

Test cases

case 1

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

expected: [[0,0],[1,-1],[2,-2]]

case 2

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

expected: [[1],[2.5],[5.5],[11.5]]

case 3

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

expected: [[1,2]]

Starter codePython

def world_model_rollout(state, weight, bias, steps):
    # Your code here
    pass

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