Environment:Facebookresearch Habitat lab CUDA GPU Training Environment

Knowledge Sources	Habitat-Lab habitat-baselines RL requirements
Domains	Infrastructure, Deep_Learning, Reinforcement_Learning
Last Updated	2026-02-15 00:00 GMT

Overview

CUDA-enabled GPU environment with PyTorch >= 1.3.1, TensorBoard, and TorchVision for RL training, policy inference, and evaluation in Habitat-Lab.

Description

This environment provides the GPU-accelerated stack required for reinforcement learning training in Habitat-Lab. It extends the core Python environment with PyTorch (GPU build), TorchVision, TensorBoard for logging, and moviepy for video generation. The device selection logic auto-detects CUDA availability and falls back to CPU. For distributed training (DD-PPO), NCCL backend and multi-GPU support are required. Optional dependencies include Weights & Biases for experiment tracking.

Usage

Use this environment for PPO/DD-PPO Training, Policy Evaluation, and Hierarchical RL Training workflows. It is the mandatory prerequisite for running `PPOTrainer.train()`, `HabitatEvaluator.evaluate_agent()`, `HierarchicalPolicy`, and all RL-based implementations.

System Requirements

Category	Requirement	Notes
OS	Linux (Ubuntu recommended)	macOS for CPU-only; Windows not supported
Hardware	NVIDIA GPU with CUDA support	CPU fallback available but much slower
VRAM	4GB+ minimum	Depends on observation resolution and model size
CUDA	CUDA toolkit compatible with PyTorch build	Version depends on PyTorch version

Dependencies

System Packages

NVIDIA GPU driver (compatible with CUDA toolkit)
CUDA toolkit (version matching PyTorch build)
`ffmpeg` (for video recording)

Python Packages

`torch` >= 1.3.1 (with CUDA support)
`torchvision` (matching torch version)
`tensorboard` >= 2.8.0
`moviepy` >= 1.0.1

Optional Packages

`wandb` (Weights & Biases for experiment tracking)

Credentials

The following environment variables are optional:

`WANDB_API_KEY`: Weights & Biases API key for experiment logging (only if using W&B integration).

Quick Install

# Install PyTorch with CUDA (adjust cu118 to your CUDA version)
pip install torch torchvision --index-url https://download.pytorch.org/whl/cu118

# Install RL training dependencies
pip install "moviepy>=1.0.1" "tensorboard>=2.8.0"

# Optional: Weights & Biases
pip install wandb

# Or install habitat-baselines from source
cd habitat-baselines && pip install -e .

Code Evidence

GPU device selection from `habitat-baselines/habitat_baselines/rl/ppo/ppo_trainer.py:905-911`:

def get_device(config: "DictConfig") -> torch.device:
    if torch.cuda.is_available():
        device = torch.device("cuda", config.habitat_baselines.torch_gpu_id)
        torch.cuda.set_device(device)
        return device
    else:
        return torch.device("cpu")

PyTorch inference mode compatibility from `habitat-baselines/habitat_baselines/utils/common.py:52-55`:

if hasattr(torch, "inference_mode"):
    inference_mode = torch.inference_mode
else:
    inference_mode = torch.no_grad

Optional W&B import from `habitat-baselines/habitat_baselines/common/tensorboard_utils.py:15-18`:

try:
    import wandb
except ImportError:
    wandb = None

RL requirements from `habitat-baselines/habitat_baselines/rl/requirements.txt`:

moviepy>=1.0.1
torch>=1.3.1
tensorboard>=2.8.0
torchvision

Common Errors

Error Message	Cause	Solution
`RuntimeError: CUDA out of memory`	Insufficient GPU VRAM	Reduce `num_environments` or observation resolution
`RuntimeError: CUDA error: no kernel image is available`	GPU compute capability mismatch	Reinstall PyTorch matching your GPU architecture
`ModuleNotFoundError: No module named 'torch'`	PyTorch not installed	Install PyTorch with CUDA support
`RuntimeError: Evaluation does not support distributed mode`	Attempting eval in distributed mode	Run evaluation on single GPU only

Compatibility Notes

CPU-only: Training works on CPU but is significantly slower. The code auto-detects and falls back.
Multi-GPU: Requires DD-PPO setup with SLURM or torch.distributed.launch. See SLURM_Distributed_Environment.
CUDA determinism: `torch.backends.cudnn.deterministic = True` is set in PPO agents for reproducibility.
TF32 precision: VER trainer enables TF32 for Ampere+ GPUs when available.

Related Pages

Page Connections

Double-click a node to navigate. Hold to expand connections.

Principle

Implementation

Heuristic

Environment