# Discount Model Search for Quality Diversity Optimization in High-Dimensional Measure Spaces

This repository is the implementation of Discount Model Search for Quality
Diversity Optimization in High-Dimensional Measure Spaces, submitted to
NeurIPS 2025.

## Installation

```bash
conda create --prefix ./env python=3.10
conda activate ./env
pip install -r requirements.txt

# Install JAX -- placed it separately since we are running both JAX and PyTorch
# simultaneously, so it is better to install JAX after so that the CUDA
# dependencies do not conflict.
pip install jax[cuda12]

# For StyleGAN3.
git clone https://github.com/NVlabs/stylegan3.git
curl -L 'https://api.ngc.nvidia.com/v2/models/org/nvidia/team/research/stylegan3/1/files?redirect=true&path=stylegan3-t-ffhqu-256x256.pkl' -o 'stylegan3-t-ffhqu-256x256.pkl'
```

If running in the LSI (Hiker) domain, download the LHQ256 dataset from
[here](https://github.com/universome/alis/blob/master/lhq.md). Place the dataset
in the `lhq_256_root/` directory, such that the directory structure looks like
this:

```
lhq_256_root
└── lhq_256
    ├── 0000000.png
    ├── 0000001.png
    ├── 0000002.png
    ├── 0000003.png
    ...
```

### Generating Centroids

We have already generated centroids and stored them in this repo under
`src/cvt/centroids/`. However, new centroids can be generated as follows. Place
any generated files under `src/cvt/centroids/`, or modify the `centroid_file`
path in the configs in `config/algo/archive/`.

1. 10D LP domains:
   ```bash
   python -m src.cvt.make_centroids \
       domain=flat_multi_100d \
       eval/archive=cvt_multi_100d
   ```
1. TA (MNIST):
   ```bash
   python -m src.cvt.mnist_centroids
   ```
1. TA (F-MNIST):
   ```bash
   python -m src.cvt.fashion_mnist_centroids
   ```
1. LSI (Hiker): Make sure to download the LHQ256 dataset as described above.
   Then:
   ```bash
   python -m src.cvt.lhq_centroids
   ```

## Experiments

The following commands will run DMS and all baselines for the paper in all
domains. Note that we use the following names for the domains in the paper:

| Name in Paper      | Name in Code                                               |
| :----------------- | :--------------------------------------------------------- |
| 2D LP (Sphere)     | `sphere_100d`                                              |
| 10D LP (Sphere)    | `sphere_multi_100d`                                        |
| 2D LP (Rastrigin)  | `rastrigin_100d`                                           |
| 10D LP (Rastrigin) | `rastrigin_multi_100d`                                     |
| 2D LP (Flat)       | `flat_100d`                                                |
| 10D LP (Flat)      | `flat_multi_100d`                                          |
| TA (MNIST)         | `triangles_mnist`                                          |
| TA (F-MNIST)       | `triangles_mnist` -- with different config from TA (MNIST) |
| LSI (Hiker)        | `lsi_face`                                                 |

```bash
# DMS -- for the benchmark domains, we include commands to run all ablations in
# the paper, i.e., learning rate, restart rule, and number of empty points.
bash scripts/paper/discount_100d.sh sphere_100d lr
bash scripts/paper/discount_100d.sh sphere_100d restart_100
bash scripts/paper/discount_100d.sh sphere_100d n_empty
bash scripts/paper/discount_multi_100d.sh sphere_multi_100d lr
bash scripts/paper/discount_multi_100d.sh sphere_multi_100d restart_basic
bash scripts/paper/discount_multi_100d.sh sphere_multi_100d n_empty
bash scripts/paper/discount_100d.sh rastrigin_100d lr
bash scripts/paper/discount_100d.sh rastrigin_100d restart_100
bash scripts/paper/discount_100d.sh rastrigin_100d n_empty
bash scripts/paper/discount_multi_100d.sh rastrigin_multi_100d lr
bash scripts/paper/discount_multi_100d.sh rastrigin_multi_100d restart_basic
bash scripts/paper/discount_multi_100d.sh rastrigin_multi_100d n_empty
bash scripts/paper/discount_100d.sh flat_100d lr
bash scripts/paper/discount_100d.sh flat_100d restart_100
bash scripts/paper/discount_100d.sh flat_100d n_empty
bash scripts/paper/discount_multi_100d.sh flat_multi_100d lr
bash scripts/paper/discount_multi_100d.sh flat_multi_100d restart_basic
bash scripts/paper/discount_multi_100d.sh flat_multi_100d n_empty
bash scripts/paper/discount_100d.sh arm_100d lr
bash scripts/paper/discount_100d.sh arm_100d restart_100
bash scripts/paper/discount_100d.sh arm_100d n_empty
bash scripts/paper/discount_triangles_mnist.sh
bash scripts/paper/discount_triangles_fashion_mnist.sh
bash scripts/paper/discount_lsi_face.sh

# CMA-MAE
bash scripts/paper/cma_mae_100d.sh sphere_100d
bash scripts/paper/cma_mae_multi_100d.sh sphere_multi_100d
bash scripts/paper/cma_mae_100d.sh rastrigin_100d
bash scripts/paper/cma_mae_multi_100d.sh rastrigin_multi_100d
bash scripts/paper/cma_mae_100d.sh flat_100d
bash scripts/paper/cma_mae_multi_100d.sh flat_multi_100d
bash scripts/paper/cma_mae_100d.sh flat_100d
bash scripts/paper/cma_mae_triangles_mnist.sh
bash scripts/paper/cma_mae_triangles_fashion_mnist.sh
bash scripts/paper/cma_mae_lsi_face.sh

# DDS
bash scripts/paper/dds_100d.sh sphere_100d
bash scripts/paper/dds_multi_100d.sh sphere_multi_100d
bash scripts/paper/dds_100d.sh rastrigin_100d
bash scripts/paper/dds_multi_100d.sh rastrigin_multi_100d
bash scripts/paper/dds_100d.sh flat_100d
bash scripts/paper/dds_multi_100d.sh flat_multi_100d
bash scripts/paper/dds_100d.sh arm_100d

# MAP-Elites (line)
bash scripts/paper/map_elites_line_100d.sh sphere_100d
bash scripts/paper/map_elites_line_multi_100d.sh sphere_multi_100d
bash scripts/paper/map_elites_line_100d.sh rastrigin_100d
bash scripts/paper/map_elites_line_multi_100d.sh rastrigin_multi_100d
bash scripts/paper/map_elites_line_100d.sh flat_100d
bash scripts/paper/map_elites_line_multi_100d.sh flat_multi_100d
bash scripts/paper/map_elites_line_100d.sh arm_100d
bash scripts/paper/map_elites_line_triangles_mnist.sh
bash scripts/paper/map_elites_line_triangles_fashion_mnist.sh
bash scripts/paper/map_elites_line_lsi_face.sh

# MAP-Elites
bash scripts/paper/map_elites_100d.sh sphere_100d
bash scripts/paper/map_elites_multi_100d.sh sphere_multi_100d
bash scripts/paper/map_elites_100d.sh rastrigin_100d
bash scripts/paper/map_elites_multi_100d.sh rastrigin_multi_100d
bash scripts/paper/map_elites_100d.sh flat_100d
bash scripts/paper/map_elites_multi_100d.sh flat_multi_100d
bash scripts/paper/map_elites_100d.sh arm_100d
bash scripts/paper/map_elites_triangles_mnist.sh
bash scripts/paper/map_elites_triangles_fashion_mnist.sh
bash scripts/paper/map_elites_lsi_face.sh
```

## Analysis

Running the experiments above will produce logging directories under the `logs`
directory, e.g., `logs/discount_arm_100d_lr/`. If desired, move the results into
a new `results/` directory, and in that directory, create a `manifest.yaml`
listing out the directories -- refer to `analysis/src/analysis/figures.py` for
the format of this manifest.

Follow the instructions in `analysis/README.md` for setting up an environment
and generating plots with the given manifest. Note that the analysis environment
is separate from the environment for the main repo.

## Results

We obtain the following results in each domain:

### 2D LP (Sphere)

|                   | QD Score             | Coverage           |
| :---------------- | :------------------- | :----------------- |
| DMS               | **6,978.20 ± 17.94** | **95.89 ± 0.43\%** |
| CMA-MAE           | 6,327.90 ± 3.60      | 80.95 ± 0.06\%     |
| DDS               | 3,156.24 ± 91.07     | 70.75 ± 0.65\%     |
| MAP-Elites (line) | 4,908.81 ± 2.48      | 60.42 ± 0.04\%     |
| MAP-Elites        | 4,163.41 ± 4.10      | 50.76 ± 0.07\%     |

### 10D LP (Sphere)

|                   | QD Score             | Coverage           |
| :---------------- | :------------------- | :----------------- |
| DMS               | **6,409.50 ± 13.50** | **89.21 ± 0.18\%** |
| CMA-MAE           | 608.53 ± 24.78       | 6.95 ± 0.31\%      |
| DDS               | 4,237.72 ± 21.18     | 60.07 ± 0.28\%     |
| MAP-Elites (line) | 2,570.74 ± 39.16     | 29.20 ± 0.45\%     |
| MAP-Elites        | 228.65 ± 8.79        | 2.35 ± 0.09\%      |

### 2D LP (Rastrigin)

|                   | QD Score             | Coverage           |
| :---------------- | :------------------- | :----------------- |
| DMS               | **5,738.90 ± 19.78** | **91.67 ± 0.44\%** |
| CMA-MAE           | 5,258.59 ± 5.97      | 80.14 ± 0.11\%     |
| DDS               | 2,495.11 ± 93.14     | 71.68 ± 0.96\%     |
| MAP-Elites (line) | 3,841.05 ± 4.67      | 56.63 ± 0.07\%     |
| MAP-Elites        | 3,172.59 ± 8.43      | 48.21 ± 0.14\%     |

### 10D LP (Rastrigin)

|                   | QD Score             | Coverage           |
| :---------------- | :------------------- | :----------------- |
| DMS               | **5,138.81 ± 10.31** | **88.19 ± 0.12\%** |
| CMA-MAE           | 246.55 ± 36.93       | 2.98 ± 0.49\%      |
| DDS               | 3,331.70 ± 18.20     | 59.54 ± 0.27\%     |
| MAP-Elites (line) | 2,001.76 ± 23.09     | 28.04 ± 0.33\%     |
| MAP-Elites        | 499.66 ± 14.06       | 7.09 ± 0.22\%      |

### 2D LP (Flat)

|                   | QD Score             | Coverage           |
| :---------------- | :------------------- | :----------------- |
| DMS               | **7,902.05 ± 35.68** | **79.02 ± 0.36\%** |
| CMA-MAE           | 5,675.90 ± 46.04     | 56.76 ± 0.46\%     |
| DDS               | 6,967.75 ± 98.80     | 69.68 ± 0.99\%     |
| MAP-Elites (line) | 4,510.65 ± 41.43     | 45.11 ± 0.41\%     |
| MAP-Elites        | 4,327.00 ± 17.29     | 43.27 ± 0.17\%     |

### 10D LP (Flat)

|                   | QD Score             | Coverage           |
| :---------------- | :------------------- | :----------------- |
| DMS               | **7,982.15 ± 24.43** | **79.82 ± 0.24\%** |
| CMA-MAE           | 1,554.90 ± 26.04     | 15.55 ± 0.26\%     |
| DDS               | 6,004.95 ± 21.52     | 60.05 ± 0.22\%     |
| MAP-Elites (line) | 757.75 ± 20.10       | 7.58 ± 0.20\%      |
| MAP-Elites        | 125.65 ± 15.88       | 1.26 ± 0.16\%      |

### Arm Repertoire

|                   | QD Score            | Coverage           |
| :---------------- | :------------------ | :----------------- |
| DMS               | **7,963.44 ± 2.47** | 80.15 ± 0.01\%     |
| CMA-MAE           | 7,902.43 ± 1.93     | 79.22 ± 0.02\%     |
| DDS               | 5,568.23 ± 54.86    | **80.24 ± 0.00\%** |
| MAP-Elites (line) | 7,458.67 ± 2.56     | 75.60 ± 0.03\%     |
| MAP-Elites        | 7,411.10 ± 2.56     | 75.42 ± 0.03\%     |

### TA (MNIST)

|                   | QD Score          | Coverage           |
| :---------------- | :---------------- | :----------------- |
| DMS               | 951.56 ± 0.95     | **99.84 ± 0.10\%** |
| CMA-MAE           | **954.27 ± 0.69** | 99.48 ± 0.09\%     |
| DDS               | --                | --                 |
| MAP-Elites (line) | 945.60 ± 1.38     | 98.86 ± 0.15\%     |
| MAP-Elites        | 941.94 ± 2.01     | 98.42 ± 0.22\%     |

### TA (F-MNIST)

|                   | QD Score          | Coverage           |
| :---------------- | :---------------- | :----------------- |
| DMS               | **701.14 ± 5.65** | **72.28 ± 0.60\%** |
| CMA-MAE           | 625.65 ± 8.60     | 63.92 ± 0.88\%     |
| DDS               | --                | --                 |
| MAP-Elites (line) | 551.13 ± 6.16     | 56.68 ± 0.64\%     |
| MAP-Elites        | 513.13 ± 6.80     | 52.68 ± 0.70\%     |

### LSI (Hiker)

|                   | QD Score              | Coverage          |
| :---------------- | :-------------------- | :---------------- |
| DMS               | **214.91 ± 15.23**    | 3.77 ± 0.13\%     |
| CMA-MAE           | 14.61 ± 26.18         | 1.56 ± 0.05\%     |
| DDS               | --                    | --                |
| MAP-Elites (line) | -51,827.44 ± 17953.11 | **7.49 ± 1.03\%** |
| MAP-Elites        | -18,917.87 ± 5705.28  | 5.06 ± 1.17\%     |

## License

This code is currently under submission for NeurIPS 2025. We will open source it
upon acceptance.
