Integrate Versatile Behavior Diffusion (VBD) model (#283)

* NYU HPC Greene setup (#316)

* Add .env for private variables and update reqs

* Add sbatch generation script

* Add env templplate

* Combine into single sbatch script

* Add logs folde

* Version changes to make pufferlib compatible

* Add jax version

* Remove old files

* rmv old files

* Minor: Fix path

* Use typer to override default args using CLI

* Add versions

* WIP

* Add optional warmup period

* Add scene selector return info

* Add VBD model and util functions

* Linting

* Add VBD model and utils

* Cleanup notebook

* Use StateDynamics model

* Refactor expert actions function

* Add minimal reproducable example

* WIP

* Rebase over main

* Warmup period steps vehicles with log playback trajectory for the first init_steps of a scene.

* Fix nans in diffusion model pred trajs

* First pass working vbd

* Small fixes

* Align road types

* Get id from waymax scenarios

* wip

* Align sims with same scenarios

* Complete align sim notebook

* wip

* Bug: Polyline dist collapsed to only 0s

* Init setup

* Add plots roadgraph

* Add more dist plots

* Update

* Rebase over main

* WIP

* Revert changes related to road graph type conversion

* Revert type additions

* Remove conversion functions

* Big resructure

* Integrate new polyline construction code

* New debug file with map_element_id and road_ids

* Merge branch 'main' of https://github.com/Emerge-Lab/gpudrive into integrations/vbd

* Use newly integrated maptypes

* Minor

* integrate new road types

* clean test notebook

* Delete old debugging scene

* Cleanup gitignore

* Feature comparison debug

* Feature comparison debug

* First pass feature analysis

* Convert GPUDrive yaw to be between [-pi, pi]

* Cleanup

* Add VBD predicted outputs to notebooks for inspection

* Fix typo in get_obs()

* revert

* refactor

* path fix

* add example

* new data

* Add debugging notebook

* Wrap yaws

* Update data structure

* Cleanup

* Update

* Decrease img sizes

* Update

* update

* regenerated debug scene

* Remove corrupted data

* Applying APs changes (todo: rebase over main after PR is merged)

* Minor

* polyline feature alignment

* minor

* Cleanup

* minor

* minor

* Integrate new datatypes

* rg

* roadgraph, deepcopy, refactoring

* roadgraph and vbd_data

* Add debugging scenes

* 50 debug scenes

* rebase over main

* rebase over main

* Visualization updates and support to train on fixed dataset size (#317)

* Add support for resample_limit

* Add option to draw the log replay trajectories in visualizer

* WIP

* minimal set maps script

* minimal set maps script

* Data processing utils

* Rendering improvements and ppo

* Make dataset to resample from configurable

* Bug fix: resampling

* Agent observation figures

* Cleanup

* Prep experiments

* Prep experiments

* vbd_data restructuring

* Fix breaking visualizer changes

* Update vbd_types

* err_val handling in vbd_data

* clone road graph tensor

* output gifs, debug scenes update

* Make collision behavior transient (#318)

* Make collision behavior transient

* Simpler way to reset collision state for collision behavior ignore

* Reset info on collision ignore

---------

Co-authored-by: Aarav Pandya <[email protected]>

* Improved visualizer and gym environment (cloning all tensors) (#320)

* Missing function

* vbd trajectory in agent obs plot

* misc fixes

* partial training integration

* Complete first 2 todo's

* Reward penalty pseudocode

* Add vbd_in_obs argument to model

* partial integration

* more partial integration changes

* make VBD work with 'good' agents

* global->local coord transform

* norm vbd obs

* discard artifacts

* move vbd to gpudrive/integrations

* moving stuff

* more deletions

* more deletions

* full cleanup

* Revert

* Improve get_obs() function

* rmv unused function

* Partial vbd updates

* minor fix

* Ensure that init_steps > 0 when using vbd

* Integrate VBD with PPO script

* Remove tutorial 9 for now

* misc

* typos

---------

Co-authored-by: Daphne Cornelisse <[email protected]>
Co-authored-by: kevin <[email protected]>
Co-authored-by: Aarav Pandya <[email protected]>

Author: 3 people

Diff for: .env.template

@@ -0,0 +1,13 @@
+# .env template
+
+# Path for logs
+LOG_FOLDER=
+
+# Your HPC account code
+NYU_HPC_ACCOUNT=
+
+# NYU ID
+USERNAME=
+
+SINGULARITY_IMAGE=
+OVERLAY_FILE=

Diff for: .gitignore

@@ -8,6 +8,7 @@
 .vscode/launch.json
 .vscode/settings.json
 .vscode/tasks.json
+
 /examples/benchmarks/results/
 /baselines/ppo/logs/*
 *.sif
@@ -24,8 +25,8 @@
 hpc/overlay*
 data/raw/*
 data/processed/validation/*
-data/processed/testing/*
 data/processed/training/*
+data/processed/testing/*
 data/processed/sampled/*
 data/processed/hand_designed/*
 analyze/figures/*

Diff for: baselines/imitation_data_generation.py

@@ -276,4 +276,4 @@ def generate_state_action_pairs(
 
     # Uncommment to save the expert actions and observations
     # torch.save(expert_actions, "expert_actions.pt")
-    # torch.save(expert_obs, "expert_obs.pt")
+    # torch.save(expert_obs, "expert_obs.pt")

Diff for: baselines/ppo/config/ppo_base_puffer.yaml

@@ -8,16 +8,16 @@ model_cpt: null
 
 environment: # Overrides default environment configs (see pygpudrive/env/config.py)
   name: "gpudrive"
-  num_worlds: 100 # Number of parallel environments
-  k_unique_scenes: 100 # Number of unique scenes to sample from
+  num_worlds: 75 # Number of parallel environments
+  k_unique_scenes: 75 # Number of unique scenes to sample from
   max_controlled_agents: 64 # Maximum number of agents controlled by the model. Make sure this aligns with the variable kMaxAgentCount in src/consts.hpp
   ego_state: true
   road_map_obs: true
   partner_obs: true
   norm_obs: true
   remove_non_vehicles: true # If false, all agents are included (vehicles, pedestrians, cyclists)
   lidar_obs: false # NOTE: Setting this to true currently turns of the other observation types
-  reward_type: "weighted_combination" 
+  reward_type: "weighted_combination"
   collision_weight: -0.75
   off_road_weight: -0.75
   goal_achieved_weight: 1.0
@@ -29,6 +29,13 @@ environment: # Overrides default environment configs (see pygpudrive/env/config.
   obs_radius: 50.0 # Visibility radius of the agents
   action_space_steer_disc: 13
   action_space_accel_disc: 7
+  # Versatile Behavior Diffusion (VBD): This will slow down training
+  use_vbd: false
+  vbd_model_path: "gpudrive/integrations/vbd/weights/epoch=18.ckpt"
+  init_steps: 11
+  vbd_trajectory_weight: 0.1 # Importance of distance to the vbd trajectories in the reward function
+  vbd_in_obs: false
+
 wandb:
   entity: ""
   project: "gpudrive"
@@ -55,7 +62,7 @@ train:
   # # # PPO # # #
   torch_deterministic: false
   total_timesteps: 1_000_000_000
-  batch_size: 262_144
+  batch_size: 131_072
   minibatch_size: 8192
   learning_rate: 3e-4
   anneal_lr: false

Diff for: baselines/ppo/ppo_pufferlib.py

@@ -170,6 +170,11 @@ def run(
     obs_radius: Annotated[Optional[float], typer.Option(help="The radius for the observation")] = None,
     collision_behavior: Annotated[Optional[str], typer.Option(help="The collision behavior; 'ignore' or 'remove'")] = None,
     remove_non_vehicles: Annotated[Optional[int], typer.Option(help="Remove non-vehicles from the scene; 0 or 1")] = None,
+    use_vbd: Annotated[Optional[bool], typer.Option(help="Use VBD model for trajectory predictions")] = False,
+    vbd_model_path: Annotated[Optional[str], typer.Option(help="Path to VBD model checkpoint")] = None,
+    vbd_trajectory_weight: Annotated[Optional[float], typer.Option(help="Weight for VBD trajectory deviation penalty")] = 0.1,
+    vbd_in_obs: Annotated[Optional[bool], typer.Option(help="Include VBD predictions in the observation")] = False,
+    init_steps: Annotated[Optional[int], typer.Option(help="Environment warmup steps")] = 0,
     # Train options
     seed: Annotated[Optional[int], typer.Option(help="The seed for training")] = None,
     learning_rate: Annotated[Optional[float], typer.Option(help="The learning rate for training")] = None,
@@ -210,10 +215,16 @@ def run(
         "remove_non_vehicles": None
         if remove_non_vehicles is None
         else bool(remove_non_vehicles),
+        "use_vbd": use_vbd,
+        "vbd_model_path": vbd_model_path,
+        "vbd_trajectory_weight": vbd_trajectory_weight,
+        "vbd_in_obs": vbd_in_obs,
+        "init_steps": init_steps,
     }
     config.environment.update(
         {k: v for k, v in env_config.items() if v is not None}
     )
+
     train_config = {
         "seed": seed,
         "learning_rate": learning_rate,

Diff for: data_utils/process_waymo_files.py

@@ -13,7 +13,8 @@
 import psutil
 from pathlib import Path
 import warnings
-from typing import Any, Dict, Optional
+from typing import Any, Dict, Optional, List
+from pdb import set_trace as T
 from tqdm import tqdm
 from waymo_open_dataset.protos import scenario_pb2, map_pb2
 from datatypes import MapElementIds
@@ -132,10 +133,10 @@ def _parse_object_state(
         "width": final_state.width,
         "length": final_state.length,
         "height": final_state.height,
-        "heading": [
-            wrap_yaws(state.heading) if state.valid else ERR_VAL
+        "heading": [ # In radians between [-pi, pi]
+            (state.heading + np.pi) % (2 * np.pi) - np.pi if state.valid else ERR_VAL
             for state in states
-        ],
+        ], 
         "velocity": [
             {"x": state.velocity_x, "y": state.velocity_y}
             if state.valid
@@ -668,7 +669,7 @@ def process_data(args):
 
     parser = argparse.ArgumentParser(
         description="Convert TFRecord files to JSON. \
-            Note: This takes about 45 seconds per tfrecord file (=50 traffic scenes)."
+            Note: This takes about 45 seconds per tfrecord file (=500 traffic scenes)."
     )
     parser.add_argument(
         "tfrecord_dir", help="Path to the directory containing TFRecord files"