from .navigation_dataset import NavigationDataset
from .procgen_dataset import ProcgenDataset
from ..simulators import SIMULATOR, Navigation, Procgen

if SIMULATOR == Navigation:
    Dataset = NavigationDataset
elif SIMULATOR == Procgen:
    Dataset = ProcgenDataset
else:
    raise NotImplementedError("Dataset not implemented for the current SIMULATOR!")
