World-Model Prediction Step

Problem

Predict the next state with a linear world model: next = W * state + b.

weight is a matrix with one row per output dimension and bias is the matching offset vector. Return the predicted vector.

Hint

This is a one-step matrix-vector prediction with a bias.

case 1

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

expected: [2,1,3]

case 2

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

expected: [5]

case 3

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

expected: [1,2]

Starter codePython

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

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