Make VBD model imports optional (#397)

* Optional vbd imports

* use_vdb = False by default

* Add optional dependencies pytorch-lightning and jaxlib

Author: daphne-cornelisse

Diff for: README.md

@@ -94,10 +94,10 @@ Finally, install the Python components of the repository using pip:
 
 ```bash
 # macOS and Linux.
-pip install -e . 
+pip install -e .
 ```
 
-Optional depencies include [pufferlib], [sb3] and [tests].
+Optional depencies include [pufferlib], [sb3], [vbd], and [tests].
 
 ```bash
 # On Windows.
@@ -113,22 +113,22 @@ pip install -e . -Cpackages.madrona_escape_room.ext-out-dir=PATH_TO_YOUR_BUILD_D
 <details>
   <summary> 🐳  Option 2. Docker </summary>
 
-To get started quickly, we provide a Dockerfile in the root directory.  
+To get started quickly, we provide a Dockerfile in the root directory.
 
-### Prerequisites  
-Ensure you have the following installed:  
-- [Docker](https://docs.docker.com/get-docker/)  
-- [NVIDIA Container Toolkit](https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/latest/install-guide.html)  
+### Prerequisites
+Ensure you have the following installed:
+- [Docker](https://docs.docker.com/get-docker/)
+- [NVIDIA Container Toolkit](https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/latest/install-guide.html)
 
-### Building the Docker mage  
-Once installed, you can build the container with:  
+### Building the Docker mage
+Once installed, you can build the container with:
 
 ```bash
 DOCKER_BUILDKIT=1 docker build --build-arg USE_CUDA=true --tag my_image:latest --progress=plain .
 ```
 
-### Running the Container  
-To run the container with GPU support and shared memory:  
+### Running the Container
+To run the container with GPU support and shared memory:
 
 ```bash
 docker run --gpus all -it --rm --shm-size=20G -v ${PWD}:/workspace my_image:latest /bin/bash
@@ -206,15 +206,15 @@ cd build
 
 ## Pre-trained policies
 
-Several pre-trained policies are available via the `PyTorchModelHubMixin` class on 🤗 huggingface_hub. 
+Several pre-trained policies are available via the `PyTorchModelHubMixin` class on 🤗 huggingface_hub.
 
-- **Best Policy (10,000 Scenarios).** The best policy from [Building reliable sim driving agents by scaling self-play](https://arxiv.org/abs/2502.14706) is available here [here](https://huggingface.co/daphne-cornelisse/policy_S10_000_02_27). This policy was trained on 10,000 randomly sampled scenarios from the WOMD training dataset. 
+- **Best Policy (10,000 Scenarios).** The best policy from [Building reliable sim driving agents by scaling self-play](https://arxiv.org/abs/2502.14706) is available here [here](https://huggingface.co/daphne-cornelisse/policy_S10_000_02_27). This policy was trained on 10,000 randomly sampled scenarios from the WOMD training dataset.
 
 - **Alternative Policy (1,000 Scenarios).** A policy trained on 1,000 scenarios can be found [here](https://huggingface.co/daphne-cornelisse/policy_S1000_02_27)
 
 ---
 
-> Note: These models were trained with the environment configurations defined in `examples/experimental/config/reliable_agents_params.yaml`, changing environment/observation configurations will affect performance. 
+> Note: These models were trained with the environment configurations defined in `examples/experimental/config/reliable_agents_params.yaml`, changing environment/observation configurations will affect performance.
 
 ---
 
@@ -240,7 +240,7 @@ See [tutorial 04](https://github.com/Emerge-Lab/gpudrive/tree/main/examples/tuto
 <details>
   <summary>Download the dataset</summary>
 
-To download the dataset you need the huggingface_hub library 
+To download the dataset you need the huggingface_hub library
 
 ```bash
 pip install huggingface_hub
@@ -356,7 +356,7 @@ and that's it!
 If you use GPUDrive in your research, please cite our ICLR 2025 paper
 ```bibtex
 @inproceedings{kazemkhani2025gpudrive,
-      title={GPUDrive: Data-driven, multi-agent driving simulation at 1 million FPS}, 
+      title={GPUDrive: Data-driven, multi-agent driving simulation at 1 million FPS},
       author={Saman Kazemkhani and Aarav Pandya and Daphne Cornelisse and Brennan Shacklett and Eugene Vinitsky},
       booktitle={Proceedings of the International Conference on Learning Representations (ICLR)},
       year={2025},

Diff for: gpudrive/env/env_torch.py

@@ -7,8 +7,6 @@
 import mediapy as media
 import gymnasium
 
-from gpudrive.integrations.vbd.data_utils import process_scenario_data
-
 from gpudrive.datatypes.observation import (
     LocalEgoState,
     GlobalEgoState,
@@ -33,8 +31,7 @@
 from gpudrive.env.dataset import SceneDataLoader
 from gpudrive.utils.geometry import normalize_min_max
 
-# Versatile Behavior Diffusion model
-from gpudrive.integrations.vbd.sim_agent.sim_actor import VBDTest
+from gpudrive.integrations.vbd.data_utils import process_scenario_data
 
 
 class GPUDriveTorchEnv(GPUDriveGymEnv):
@@ -132,12 +129,12 @@ def __init__(
 
     def _initialize_vbd(self):
         """
-        Initialize the Versatile Behavior Diffusion (VBD) model and related components.
+        Initialize the Versatile Behavior Diffusion (VBD) model and related
+        components. Link: https://arxiv.org/abs/2404.02524.
 
         Args:
             config: Configuration object containing VBD settings.
         """
-        # Set VBD configuration parameters
         self.use_vbd = self.config.use_vbd
         self.vbd_trajectory_weight = self.config.vbd_trajectory_weight
 
@@ -149,16 +146,13 @@ def _initialize_vbd(self):
         else:
             self.init_steps = self.config.init_steps
 
-        # Initialize VBD model and trajectories if enabled
         if (
             self.use_vbd
             and hasattr(self.config, "vbd_model_path")
             and self.config.vbd_model_path
         ):
-            # Load VBD model from specified path
             self.vbd_model = self._load_vbd_model(self.config.vbd_model_path)
 
-            # Initialize trajectory tensor with zeros
             self.vbd_trajectories = torch.zeros(
                 (
                     self.num_worlds,
@@ -167,17 +161,18 @@ def _initialize_vbd(self):
                     5,
                 ),
                 device=self.device,
-                dtype=torch.float32,  # Explicitly specify dtype for clarity
+                dtype=torch.float32,
             )
 
             self._generate_vbd_trajectories()
         else:
-            # Set to None if VBD is disabled or model path is not provided
             self.vbd_model = None
             self.vbd_trajectories = None
 
     def _load_vbd_model(self, model_path):
         """Load the Versatile Behavior Diffusion (VBD) model from checkpoint."""
+        from gpudrive.integrations.vbd.sim_agent.sim_actor import VBDTest
+
         model = VBDTest.load_from_checkpoint(
             model_path, torch.device(self.device)
         )
@@ -1169,6 +1164,7 @@ def get_obs(self, mask=None):
         ego_states = self._get_ego_state(mask)
         partner_observations = self._get_partner_obs(mask)
         road_map_observations = self._get_road_map_obs(mask)
+
         if (
             self.use_vbd
             and self.vbd_model is not None

Diff for: pyproject.toml

@@ -50,6 +50,12 @@ test = [
     "pytest>=8.2.1",
 ]
 
+vbd = [
+    "lightning",
+    "jaxlib",
+    "waymo-waymax @ git+https://github.com/waymo-research/waymax.git@main",
+]
+
 [tool.madrona.packages.madrona_gpudrive]
 ext-only = true
 ext-out-dir = "build"