"""PlaNet RSSM rollout interface."""


def rollout(model, batch, horizon):
    images, actions, future_actions = batch
    return model.rollout(images, actions, future_actions[:, :horizon])