FlowMo-WM / experiments /flowmo /src /predict.py
cccat6's picture
Initial FlowMo-WM public code release
604e535 verified
raw
history blame
228 Bytes
"""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])