This supplementary materials contains our code, and configuration files that enable one to replicate all of our experiments and figures exactly as they appear in the paper.

The code for the paper is contained in the `HetEnvDesign` folder.

The code in the `vmas,torchrl,tensordict` folders is not part of the submission and is just needed as a dependency. These folders may contain references to the original authors and should not be reviewed.

The exception is the custom scenarios created in the folder `vmas/scenarios`, which are part of this submission.

## Installing

1. Create a virtual environment with python 3.11 (e.g., `conda create --name env python=3.11` with conda)
2. Install dependencies
```bash
pip install torch==2.5 hydra-core torch_geometric wandb moviepy matplotlib==3.8
```
3. Install our versions of `VMAS`, `TensorDict`, and `TorchRL`.

```bash
cd tensordict
python setup.py develop
cd ..

cd torchrl
python setup.py develop
cd ..

cd vmas
pip install -e .
cd ..
```

4. Try running a script (if you have a gpu, it will use it)
```bash
python HetEnvDesign/matrix_game_cont.py
```

## Running

Each experiment in the paper has a corresponding python file and yaml configuration folder (under the `conf` folder), containing its hyper-parameters.

Here is how to reproduce the experiments (any parameters in the corresponding folder can be passed).

`-m` performs a multirun, see [the hydra docs](https://hydra.cc/docs/intro/) for more info.

### One-shot Matrix Game

This is the ``matrix game'' considered in our experiments section.

#### Running
Continuous:

```bash
python HetEnvDesign/matrix_game_cont.py -m env.scenario.gen_agg_type_task=min,mean,max env.scenario.gen_agg_type_agent=min,mean,max seed=0,1,2,3,4,5,6,7,8
```

Discrete:

```bash
python HetEnvDesign/matrix_game_disc.py -m env.scenario.gen_agg_type_task=min,mean,max env.scenario.gen_agg_type_agent=min,mean,max seed=0,1,2,3,4,5,6,7,8
```
#### Plotting
Once run, you can plot the results from wandb (it will create the pdf used in the paper)
```bash
python HetEnvDesign/plotting/plot_matrix_games.py
```

### Multi-Goal Capture

This is the embodied environment studied in our experiments section.

#### Running

```bash
python HetEnvDesign/goal_navigation_vanilla.py -m env.scenario.gen_agg_type_task=min,mean,max env.scenario.gen_agg_type_agent=min,mean,max seed=0,1,2,3,4,5,6,7,8
```

#### Plotting
Once run, you can plot the results from wandb (it will create the pdf used in the paper)
```bash
python HetEnvDesign/plotting/plot_ctf_embodied_vanilla.py
```

### Football and Tag

Can be run similar to the above, using the files:

```bash
HetEnvDesign/soccer.py
HetEnvDesign/tag_potential.py
```

### Multi-Goal Capture with Local Sensing

This runs the experiments where agents are equipped with a range sensor of increasing radius.

#### Running

```bash
python HetEnvDesign/goal_navigation_lidar.py -m env.scenario.lidar_range=0,0.1,0.2,0.35 seed=0,1,2,3
```

#### Plotting
Once run, you can plot the results from wandb (it will create the pdf used in the paper)
```bash
python HetEnvDesign/plotting/plot_ctf_embodied_lidar.py
```

### Heterogeneous Environment Design (HED) experiments

These are the co-design HED experiments where HED optimizes the reward structure to favour heterogeneity.

#### Running

Softmax:
```bash
python HetEnvDesign/goal_navigation_hed_softmax.py -m seed=0,1,2,3,4,5,6,7,8,9,10,11,12
```

Power-Sum:
```bash
python HetEnvDesign/goal_navigation_hed_powersum.py -m seed=0,1,2,3,4,5,6,7,8,9,10,11,12
```

#### Plotting
Once run, you can plot the results from wandb (it will create the pdf used in the paper)
```bash
python HetEnvDesign/plotting/plot_ctf_embodied_softmax_design.py
python HetEnvDesign/plotting/plot_ctf_embodied_powersum_design.py
```

## Configuring

The configuration for the codebase is available in the [`HetEnvDesign/conf`](het_control/conf) folder.

Each parameter can be changed from the yaml file or from command line using the hydra syntax, see [the hydra docs](https://hydra.cc/docs/intro/) for more info.
