Dream to Drive with Predictive Individual World Model

JAX implementation of the PIWM. PIWM is a novel model-based reinforcement learning method built upon DreamerV3. It is typically designed for autonomous driving tasks. It learns to model the driving environment from an individual perspective and enhances the transition dynamic by explicitly modeling the interactive relations between vehicles. Trajectory prediction further replaces observation reconstruction for representation learning to better capture the future intentions or motion trends of interested vehicles within the latent states.

Project page

The simulation environment used in this work is I-SIM simulator, which is built upon Interaction Dataset.

• Overall Framework and Differences from Original Dreamer:

• Detailed Structure of Models:

Instructions

The code has been tested on Linux and requires Python 3.11+.

Docker

You can either use the provided Dockerfile that contains instructions or follow the manual instructions below.

Manual

Install JAX and then the other dependencies:

pip install -U -r requirements.txt

Modify the "loader_type" in dreamerv3/configs.yaml to select training or testing for large-scale or small-scale scenarios:

Large scale scenarios:

loader_type: large_scale

Small scale scenarios:

loader_type: small_scale

Modify the "script in dreamerv3/configs.yaml to select training:

"configs.yaml":

script: train

Training script:

python train.py --logdir {logdir} --configs interaction_prediction --task interaction_prediction

Modify the "script in dreamerv3/configs.yaml and specify the checkpoint file to select testing:

"configs.yaml":

script: eval_only
...
from_checkpoint: {checkpoint_dir}/{checkpoint}.ckpt

Test scripts:

python train.py --logdir {logdir} --configs interaction_prediction --task interaction_prediction

Catalog

• Code & Checkpoints Release
• Initialization

Citation

If you find PIWM useful in your research or applications, please consider giving us a star 🌟 and citing it by the following BibTeX entry.

@article{gao2024piwm,
  author={Gao, Yinfeng and Zhang, Qichao and Ding, Da-Wei and Zhao, Dongbin},
  journal={IEEE Transactions on Intelligent Vehicles}, 
  title={Dream to Drive With Predictive Individual World Model}, 
  year={2024},
  volume={},
  number={},
  pages={1-16},
  doi={10.1109/TIV.2024.3408830}}

Acknowledgement

We appreciate the following GitHub repos for their valuable code base or dataset:

https://github.com/danijar/dreamerv3