"""FlowMo rollout interface."""


def rollout(model, batch, horizon):
    """Run open-loop FlowMo prediction."""
    images, actions, future_actions = batch
    return model.rollout(images, actions, future_actions[:, :horizon])