# Learning to Drive with Two Minds: A Competitive Dual-Policy Approach in Latent World Models

This repository provides the implementation of our NeurIPS 2025 submission, built on top of [SSR](https://github.com/PeidongLi/SSR) and [WoTE](https://github.com/liyingyanUCAS/WoTE). We extend it by integrating a learned latent world model and designing a hybrid imitation learning (IL) and reinforcement learning (RL) framework.


## 🔧 Installation

Follow the instructions in `SSR/SSR/README_SSR.md` and `Wote\README_WoTE.md` to configure the Conda environment and install dependencies.


## 🚗 Our Method on NuScenes

To train and evaluate our model on the NuScenes dataset, run the following script:

```bash
bash SSR/SSR/tools/single_experiment.sh
```

This script includes training and evaluation logic. Make sure your dataset path and config settings are properly set in the script or corresponding config files.

**Arguments:**

- `--output_dir`: Output directory to save the ckpt, tensorboard information and message durting training.
- `--model_identifier`: The model identifier of this trial.
- `--config_path`: Path of configuration you want to use.


We also provide a example log file (not our final version of model~) for reviewers: `SSR\log\20250510_201154.log`.

## 🧭 Our Method on Navsim

### 1. Download the Navsim Dataset

You need to download the following parts:

- `maps`
- `navtrain`
- `test`

Use the shell scripts in the `download` folder:

```bash
bash download/download_maps.sh
bash download/download_navtrain.sh
bash download/download_test.sh
```

After downloading, organize the dataset as follows:

```
navsim_dataset/
├── maps/
│   ├── .mapblocks/
│   ├── sg-one-north/
│   ├── us-ma-boston/
│   ├── us-nv-las-vegas-strip/
│   ├── us-pa-pittsburgh-hazelwood/
│   └── nuplan-maps-v1.0.json
├── navsim_logs/
│   ├── test/
│   └── trainval/
├── sensor_blobs/
│   ├── test/
│   └── trainval/
```

---

### 2. Install Navsim Environment

```bash
git clone https://github.com/wote/navsim.git
cd navsim
conda env create -f environment.yml
conda activate wote
pip install -r requirements.txt
pip install git+https://github.com/motional/nuplan-devkit.git@nuplan-devkit-v1.2#egg=nuplan-devkit
pip install -e .
```


### 3. Training and Evaluation

#### (1) Metric Caching on Test Set

Update the dataset path in the script and run:

```bash
bash scripts/evaluation/run_metric_caching.sh
```

#### (2) Training

Update the paths for `trajectory_anchors_256.npy` and `formatted_pdm_score_256.npy` in the config file, then run:

```bash
bash scripts/training/run_wote2.sh
```

#### (3) Inference

Update the checkpoint path (make sure it matches the training version) and run:

```bash
bash scripts/training/eval_wote2.sh
```


## 🎥 Visualization Tool

We provide several visualization scripts to analyze results on NuScenes.

### 1. Trajectory and BEV Feature Attention Visualization

```bash
sh vis_gt_and_bevfeature.sh
```

**Arguments:**

- `--data_directory`: Path to the visualization input data.
- `--save_directory`: Output directory for the visualization results.


### 2. World Model Closed-Loop Exploration Visualization

```bash
sh vis_world_model_closed_loop.sh
```

**Arguments:**

- `--data_directory`: Input data path.
- `--save_directory`: Directory to save visualizations.
- `--scene_token`: Scene token to visualize.
- `--frame_idx`: Frame index to visualize.


### 3. SSR/VAD Multi-Trajectory Comparison Visualization

```bash
sh vis.sh
```

**Arguments:**

- `--vad-result-path`: Path to VAD result file.
- `--ssr-result-path`: Path to SSR result file.
- `--save_path`: Path to save visualization outputs.


## 📌 Notes

- Our method is built on [SSR](https://github.com/PeidongLi/SSR).
- Please ensure dataset paths and configs are correctly set before running any scripts.
- This codebase includes imitation learning, reinforcement learning, and a learned world model integrated into the SSR pipeline.


## 📢 Code Release

We will release our full training logs and pretrained models soon. Stay tuned!