Implementation:Facebookresearch Habitat lab EQA CNN Pretrain Trainer

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

Overview

EQACNNPretrainTrainer is a trainer class for pre-training the MultitaskCNN encoder-decoder used as a feature extractor in the Embodied Question Answering (EQA) pipeline.

Description

The EQACNNPretrainTrainer class extends BaseILTrainer and is registered under the name "eqa-cnn-pretrain" in the baseline registry. It trains the MultitaskCNN model on three simultaneous tasks: semantic segmentation, depth prediction, and autoencoder RGB reconstruction. The combined loss function is:

loss = seg_loss + 10 * ae_loss + 10 * depth_loss

The trainer supports both training and evaluation phases:

• Training: Loads the EQACNNPretrainDataset, creates a DataLoader with configurable batch size, initializes the MultitaskCNN model, and trains using Adam optimizer. Losses are logged to TensorBoard, and model checkpoints are saved at every epoch.

• Evaluation: Loads a checkpoint, runs the model on the validation set, computes average losses for all three tasks, and optionally saves visual results (RGB, segmentation, and depth reconstructions) at configurable intervals.

The trainer uses CrossEntropyLoss for segmentation and SmoothL1Loss for both depth and autoencoder reconstruction tasks.

Usage

Use this trainer to pre-train the CNN feature extractor before training the downstream VQA or navigation models. The pretrained encoder checkpoint is subsequently loaded by VqaLstmCnnAttentionModel and NavPlannerControllerModel for feature extraction.

Code Reference

Source Location

• Repository: Facebookresearch_Habitat_lab
• File: habitat-baselines/habitat_baselines/il/trainers/eqa_cnn_pretrain_trainer.py
• Lines: 1-303

Signature

@baseline_registry.register_trainer(name="eqa-cnn-pretrain")
class EQACNNPretrainTrainer(BaseILTrainer):
    supported_tasks = ["EQA-v0"]

    def __init__(self, config=None): ...
    def _make_results_dir(self) -> None: ...
    def _save_results(
        self,
        gt_rgb: torch.Tensor,
        pred_rgb: torch.Tensor,
        gt_seg: torch.Tensor,
        pred_seg: torch.Tensor,
        gt_depth: torch.Tensor,
        pred_depth: torch.Tensor,
        path: str,
    ) -> None: ...
    def train(self) -> None: ...
    def _eval_checkpoint(
        self,
        checkpoint_path: str,
        writer: TensorboardWriter,
        checkpoint_index: int = 0,
    ) -> None: ...

Import

from habitat_baselines.il.trainers.eqa_cnn_pretrain_trainer import EQACNNPretrainTrainer

I/O Contract

Inputs

Name	Type	Required	Description
config	DictConfig	Yes	Habitat baselines configuration containing dataset paths, batch size, learning rate, max epochs, and logging settings

Outputs

Name	Type	Description
checkpoint files	file	Model state dict saved as epoch_{N}.ckpt files in the configured checkpoint directory
TensorBoard logs	file	Training and evaluation loss metrics logged via TensorboardWriter
visual results	file	Optional RGB, segmentation, and depth reconstruction images saved during evaluation

Usage Examples

Basic Usage

from habitat_baselines.common.baseline_registry import baseline_registry
from habitat.config.default_structured_configs import register_hydra_plugin
from omegaconf import DictConfig

# The trainer is registered and typically invoked via the run script:
# python -m habitat_baselines.run --config-name=eqa/il_eqa_cnn_pretrain.yaml

# Programmatic usage:
config = ...  # load Hydra config
trainer_cls = baseline_registry.get_trainer("eqa-cnn-pretrain")
trainer = trainer_cls(config=config)

# Training
trainer.train()

# Evaluation (called internally by eval loop)
# trainer.eval()