# Adversarial environment design via regret guided diffusion model
This code is the implementation of ADD: adversarial environment design via regret-guided diffusion models. We implement the proposed algorithm based on the [Dual Curriculum Design (DCD)](https://github.com/facebookresearch/dcd). Also, our implementation of the diffusion model is based on the [Censored Sampling of Diffusion Models Using 3 Minutes of Human Feedback](https://github.com/tetrzim/diffusion-human-feedback).

## Setup
To install the necessary dependencies, run the following commands:
```
conda env create -f enviroment.yaml
conda activate add
pip install -r requirements.txt
git clone https://github.com/openai/baselines.git
cd baselines
pip install -e .
cd ..
pip install pyglet==1.5.11
```
Ignore error messages regarding dependecies. But you might need to install additional packages (ex. six, xvfb)

## Training a diffusion model
First, we should generate a random environment dataset to train a diffusion model. please replace {ENV_NAME} to minigrid or bipedal.
```
cd add/diffusion_human_feedback
python datasets/{ENV_NAME}.py
```
For the partially observable navigation task, run the following commands to train a diffusion model:
```
MODEL_FLAGS="--image_size 16 --image_channels 3 --num_channels 128 --num_res_blocks 3"
DIFFUSION_FLAGS="--diffusion_steps 1000 --noise_schedule linear"
TRAIN_FLAGS="--lr 1e-4 --batch_size 256 --save_interval 100000"
LOG_DIR="log/minigrid_60_uniform" 
NUM_GPUS="1" 

echo $(mpiexec -n $NUM_GPUS python image_train.py --log_dir=$LOG_DIR --data_dir=minigrid_60_uniform --rgb=True --random_flip=False $MODEL_FLAGS $DIFFUSION_FLAGS $TRAIN_FLAGS)
```
For the 2D bipedal locomotion task, run the following commands to train a diffusion model:
```
python flat_train.py
```
## Making shell commands
To make a shell commands, you must designate the log directory. 

For the partially observable navigation task, run the following commands:
```
python train_scripts/make_cmd.py --json minigrid/60_blocks_uniform/mg_60b_uni_add_categorical --num_trials 5
```
For the 2D bipedal locomotion task, run the following commands:
```
python train_scripts/make_cmd.py --json bipedal/bipedal_add --num_trials 5
```

## Run experiments
Finally, run the following commands to start the training:
```
chmod +x run.sh
sh run.sh
```

