Implementation:Facebookresearch Habitat lab Dataset

Knowledge Sources	Facebookresearch_Habitat_lab
Domains	Embodied_AI, Data_Management, Episode_Management
Last Updated	2026-02-15 00:00 GMT

Overview

Core dataset module that defines the Episode and Dataset base classes along with an EpisodeIterator for managing collections of embodied AI episodes with support for scene grouping, shuffling, cycling, and splitting.

Description

This module provides the foundational data abstractions for Habitat's task system:

BaseEpisode:

• Minimal episode specification with episode_id and scene_id.
• Used for lightweight episode identification in evaluation pipelines.

Episode (extends BaseEpisode):

• Full episode specification including start_position (3D coordinates), start_rotation (quaternion), scene_dataset_config, additional_obj_config_paths, and optional info dictionary.
• Implements shortest path cache invalidation hooks.
• Custom __getstate__/__setstate__ for efficient serialization (excludes cache).

Dataset (Generic[T]):

• Base class holding a list of episodes.
• Properties: num_episodes, scene_ids (unique sorted scene IDs).
• Querying: get_scene_episodes, get_episodes (by index).
• Iteration: get_episode_iterator returns a configurable EpisodeIterator.
• Filtering: filter_episodes creates a new dataset with filtered episodes; build_content_scenes_filter creates a filter from config.
• Splitting: get_splits divides the dataset into subsets for multi-worker training, with options for scene collation, sorting, uneven splits, and memory optimization.
• Serialization: to_json and from_json.

EpisodeIterator:

• Configurable iterator with support for:
  • Cycling: Loop back to start when all episodes are consumed.
  • Shuffling: Randomize episode order on each cycle.
  • Scene grouping: Keep episodes from the same scene together to reduce simulator overhead.
  • Max scene repeat: Limit consecutive episodes or steps from the same scene to prevent overfitting.
  • Sampling: Select a random subset of episodes.
  • Step repetition range: Randomize the scene switch threshold to prevent synchronized worker transitions.
• Methods: set_next_episode_by_index, set_next_episode_by_id, step_taken.

Usage

Subclass Dataset for specific task datasets (e.g., PointNav, ObjectNav, Rearrangement). The EpisodeIterator is used internally by the Habitat environment to iterate through episodes during training and evaluation. The splitting functionality is used for distributing work across multiple training workers.

Code Reference

Source Location

• Repository: Facebookresearch_Habitat_lab
• File: habitat-lab/habitat/core/dataset.py
• Lines: 1-584

Signature

@attr.s(auto_attribs=True, kw_only=True)
class Episode(BaseEpisode):
    scene_dataset_config: str = "default"
    additional_obj_config_paths: List[str] = []
    start_position: List[float] = None
    start_rotation: List[float] = None
    info: Optional[Dict[str, Any]] = None

class Dataset(Generic[T]):
    episodes: List[T]

    @property
    def num_episodes(self) -> int: ...
    @property
    def scene_ids(self) -> List[str]: ...
    def get_scene_episodes(self, scene_id: str) -> List[T]: ...
    def get_episodes(self, indexes: List[int]) -> List[T]: ...
    def get_episode_iterator(self, *args, **kwargs) -> Iterator[T]: ...
    def filter_episodes(self, filter_fn: Callable[[T], bool]) -> "Dataset": ...
    def get_splits(
        self,
        num_splits: int,
        episodes_per_split: Optional[int] = None,
        remove_unused_episodes: bool = False,
        collate_scene_ids: bool = True,
        sort_by_episode_id: bool = False,
        allow_uneven_splits: bool = False,
    ) -> List["Dataset"]: ...
    def to_json(self) -> str: ...
    def from_json(self, json_str: str, scenes_dir: Optional[str] = None) -> None: ...

class EpisodeIterator(Iterator[T]):
    def __init__(
        self,
        episodes: Sequence[T],
        cycle: bool = True,
        shuffle: bool = False,
        group_by_scene: bool = True,
        max_scene_repeat_episodes: int = -1,
        max_scene_repeat_steps: int = -1,
        num_episode_sample: int = -1,
        step_repetition_range: float = 0.2,
        seed: int = None,
    ) -> None: ...
    def __next__(self) -> Episode: ...
    def set_next_episode_by_index(self, episode_index: int) -> None: ...
    def set_next_episode_by_id(self, episode_id: str) -> None: ...
    def step_taken(self) -> None: ...

Import

from habitat.core.dataset import Dataset, Episode, BaseEpisode, EpisodeIterator, ALL_SCENES_MASK

I/O Contract

Inputs (Episode)

Name	Type	Required	Description
episode_id	str	Yes	Unique identifier for the episode
scene_id	str	Yes	Identifier of the scene in the dataset
start_position	List[float]	Yes	3D cartesian coordinates [x, y, z] for initial agent position
start_rotation	List[float]	Yes	Unit quaternion [x, y, z, w] for initial agent orientation
scene_dataset_config	str	No	Path to the SceneDataset config file (default "default")
additional_obj_config_paths	List[str]	No	Additional paths for object config files
info	Optional[Dict[str, Any]]	No	Additional episode metadata

Outputs (Dataset.get_splits)

Name	Type	Description
return	List[Dataset]	List of new Dataset objects, each containing a disjoint subset of episodes

Usage Examples

Filtering Episodes by Scene

from habitat.core.dataset import Dataset

# Filter to only include episodes from a specific scene
filtered_dataset = dataset.filter_episodes(
    lambda ep: "kitchen" in ep.scene_id
)
print(f"Filtered to {filtered_dataset.num_episodes} kitchen episodes")

Splitting a Dataset for Multi-Worker Training

# Split dataset into 4 worker-specific subsets
splits = dataset.get_splits(
    num_splits=4,
    collate_scene_ids=True,
    sort_by_episode_id=True,
)
for i, split in enumerate(splits):
    print(f"Worker {i}: {split.num_episodes} episodes")

Using the Episode Iterator

from habitat.core.dataset import EpisodeIterator

iterator = EpisodeIterator(
    episodes=dataset.episodes,
    cycle=True,
    shuffle=True,
    group_by_scene=True,
    max_scene_repeat_episodes=100,
    seed=42,
)

# Iterate through episodes
for episode in iterator:
    print(f"Episode {episode.episode_id} in scene {episode.scene_id}")
    iterator.step_taken()  # track steps for scene switching