Implementation:Facebookresearch Habitat lab EqaCnnPretrainDataset

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

Overview

EQACNNPretrainDataset is a PyTorch Dataset for pretraining the CNN feature extractor used in Embodied Question Answering, storing RGB, depth, and semantic segmentation frames in an LMDB database.

Description

EQACNNPretrainDataset extends torch.utils.data.Dataset and manages a cache of rendered scene observations in LMDB format. On first use, it creates a Habitat environment, iterates over all episodes grouped by scene, randomly samples 9 positions from each episode's shortest path, and renders RGB, depth, and semantic segmentation observations at those positions. The semantic segmentation is computed by mapping instance IDs to category labels. All three modalities are stored as binary blobs in LMDB. On subsequent runs, it detects the existing cache and opens it in read-only mode. The __getitem__ method lazily opens the LMDB connection and retrieves the RGB (normalized float32 CHW), depth (float32 1HW), and segmentation (uint8 HW) tensors for a given index.

Usage

Use this dataset to pretrain a CNN feature extractor on visual observations from Habitat environments before using the features for downstream EQA tasks.

Code Reference

Source Location

• Repository: Facebookresearch_Habitat_lab
• File: habitat-baselines/habitat_baselines/il/data/eqa_cnn_pretrain_data.py
• Lines: 19-179

Signature

class EQACNNPretrainDataset(Dataset):
    def __init__(self, config, mode="train"):

Import

from habitat_baselines.il.data.eqa_cnn_pretrain_data import EQACNNPretrainDataset

I/O Contract

Inputs

Name	Type	Required	Description
config	DictConfig	Yes	Configuration object containing habitat environment and baselines settings
mode	str	No	Dataset split mode, either "train" or "val" (default: "train")

Outputs

Name	Type	Description
idx	int	The sample index
rgb	np.ndarray (float32, 3x256x256)	Normalized RGB image in CHW format
depth	np.ndarray (float32, 1x256x256)	Depth map in CHW format
seg	np.ndarray (uint8, 256x256)	Semantic segmentation labels in HW format

Key Methods

save_frames

def save_frames(self, pos_queue: List[ShortestPathPoint]) -> None

Renders RGB, depth, and segmentation at each position and writes them to LMDB.

cache_exists

def cache_exists(self) -> bool

Checks whether the LMDB cache directory exists and is non-empty.

__getitem__

def __getitem__(self, idx: int) -> Tuple[int, np.ndarray, np.ndarray, np.ndarray]

Returns (index, rgb, depth, segmentation) for a given sample index.

Usage Examples

Basic Usage

from habitat_baselines.il.data.eqa_cnn_pretrain_data import EQACNNPretrainDataset
from torch.utils.data import DataLoader

train_dataset = EQACNNPretrainDataset(config=my_config, mode="train")
val_dataset = EQACNNPretrainDataset(config=my_config, mode="val")

train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)

for idx, rgb, depth, seg in train_loader:
    # rgb: (B, 3, 256, 256) float32
    # depth: (B, 1, 256, 256) float32
    # seg: (B, 256, 256) uint8
    features = cnn_model(rgb)
    loss = compute_loss(features, seg)
    loss.backward()