class RluDmlabExploreObjectRewardsFew(dmlab_dataset.DMLabDatasetBuilder):
  """DatasetBuilder for rlu_dmlab_explore_object_rewards_few dataset."""

  VERSION = tfds.core.Version('1.0.0')
  RELEASE_NOTES = {
      '1.0.0': 'Initial release.',
  }

  # pytype: disable=wrong-keyword-args
  BUILDER_CONFIGS = [
      dmlab_dataset.BuilderConfig(
          name='training_0', task=_TASK, episode_length=_EPISODE_LENGTH),
      dmlab_dataset.BuilderConfig(
          name='training_1', task=_TASK, episode_length=_EPISODE_LENGTH),
      dmlab_dataset.BuilderConfig(
          name='training_2', task=_TASK, episode_length=_EPISODE_LENGTH)
  ]
class RluDmlabRoomsSelectNonmatchingObject(dmlab_dataset.DMLabDatasetBuilder):
  """DatasetBuilder for rlu_dmlab_rooms_select_nonmatching_object dataset."""

  VERSION = tfds.core.Version('1.2.0')
  RELEASE_NOTES = {
      '1.0.0': 'Initial release.',
      '1.1.0': 'Added is_last.',
      '1.2.0': 'BGR -> RGB fix for pixel observations.',
  }

  # pytype: disable=wrong-keyword-args
  BUILDER_CONFIGS = [
      dmlab_dataset.BuilderConfig(
          name='training_0', task=_TASK, episode_length=_EPISODE_LENGTH),
      dmlab_dataset.BuilderConfig(
          name='training_1', task=_TASK, episode_length=_EPISODE_LENGTH),
      dmlab_dataset.BuilderConfig(
          name='training_2', task=_TASK, episode_length=_EPISODE_LENGTH)
  ]
class RluDmlabSeekavoidArena01(dmlab_dataset.DMLabDatasetBuilder):
    """DatasetBuilder for rlu_dmlab_seekavoid_arena01 dataset."""

    VERSION = tfds.core.Version('1.1.0')
    RELEASE_NOTES = {
        '1.0.0': 'Initial release.',
        '1.1.0': 'Added is_last.',
    }

    # pytype: disable=wrong-keyword-args
    BUILDER_CONFIGS = [
        dmlab_dataset.BuilderConfig(name=name,
                                    task=_TASK,
                                    episode_length=_EPISODE_LENGTH)
        for name in _CONFIG_NAMES
    ]
class RluDmlabSeekavoidArena01(dmlab_dataset.DMLabDatasetBuilder):
    """DatasetBuilder for rlu_dmlab_seekavoid_arena01 dataset."""

    VERSION = tfds.core.Version('1.0.0')
    RELEASE_NOTES = {
        '1.0.0': 'Initial release.',
    }
    _TASK = 'seekavoid_arena_01'
    _EPISODE_LENGTH = 301

    # pytype: disable=wrong-keyword-args
    BUILDER_CONFIGS = [
        dmlab_dataset.BuilderConfig(name='training_0',
                                    task=_TASK,
                                    episode_length=_EPISODE_LENGTH),
        # TODO(b/188125902): Enable the rest of the configs once we double check
        # that the first one works.
        #  dmlab_dataset.TaskConfig(name='training_1',
        #                         task=_TASK,
        #                         episode_length=_EPISODE_LENGTH),
        # dmlab_dataset.TaskConfig(name='training_2',
        #                         task=_TASK,
        #                         episode_length=_EPISODE_LENGTH),
        #  dmlab_dataset.TaskConfig(name='snapshot_0_eps_0.0',
        #                         task=_TASK,
        #                         episode_length=_EPISODE_LENGTH),
        # dmlab_dataset.TaskConfig(name='snapshot_1_eps_0.0',
        #                         task=_TASK,
        #                         episode_length=_EPISODE_LENGTH),
        #  dmlab_dataset.TaskConfig(name='snapshot_0_eps_0.01',
        #                         task=_TASK,
        #                         episode_length=_EPISODE_LENGTH),
        # dmlab_dataset.TaskConfig(name='snapshot_1_eps_0.01',
        #                         task=_TASK,
        #                         episode_length=_EPISODE_LENGTH),
        #  dmlab_dataset.TaskConfig(name='snapshot_0_eps_0.25',
        #                         task=_TASK,
        #                         episode_length=_EPISODE_LENGTH),
        # dmlab_dataset.TaskConfig(name='snapshot_1_eps_0.25',
        #                         task=_TASK,
        #                         episode_length=_EPISODE_LENGTH),
    ]