├── .gitignore ├── LICENSE ├── README.md ├── create_slurm_scripts.py ├── design_baselines ├── __init__.py ├── autofocused_cbas │ ├── __init__.py │ ├── distribution_experiments.py │ ├── experiments.py │ ├── nets.py │ ├── relabel_experiments.py │ └── trainers.py ├── bo_qei │ ├── __init__.py │ ├── distribution_experiments.py │ ├── experiments.py │ ├── nets.py │ ├── relabel_experiments.py │ └── trainers.py ├── cbas │ ├── __init__.py │ ├── distribution_experiments.py │ ├── experiments.py │ ├── nets.py │ ├── relabel_experiments.py │ └── trainers.py ├── cli.py ├── cma_es │ ├── __init__.py │ ├── distribution_experiments.py │ ├── experiments.py │ ├── nets.py │ ├── relabel_experiments.py │ └── trainers.py ├── coms_cleaned │ ├── __init__.py │ ├── experiments.py │ ├── nets.py │ └── trainers.py ├── coms_original │ ├── __init__.py │ ├── experiments.py │ ├── nets.py │ └── trainers.py ├── data.py ├── gradient_ascent │ ├── __init__.py │ ├── distribution_experiments.py │ ├── distribution_mean_ensemble_experiments.py │ ├── distribution_min_ensemble_experiments.py │ ├── experiments.py │ ├── mean_ensemble_experiments.py │ ├── min_ensemble_experiments.py │ ├── nets.py │ ├── relabel_experiments.py │ ├── relabel_mean_ensemble_experiments.py │ ├── relabel_min_ensemble_experiments.py │ └── trainers.py ├── logger.py ├── mins │ ├── __init__.py │ ├── distribution_experiments.py │ ├── experiments.py │ ├── nets.py │ ├── relabel_experiments.py │ ├── replay_buffer.py │ ├── trainers.py │ └── utils.py ├── reinforce │ ├── __init__.py │ ├── distribution_experiments.py │ ├── experiments.py │ ├── nets.py │ ├── online_experiments.py │ ├── relabel_experiments.py │ └── trainers.py └── utils.py ├── environment.yml ├── fine-tuning ├── coms-chembl.sh ├── coms-nas.sh ├── coms-tf-bind-10.sh └── coms-tf-bind-8.sh ├── label_with_scores.py ├── process_nas_results.py ├── requirements.txt ├── scripts ├── coms-ant.sh ├── coms-chembl.sh ├── coms-dkitty.sh ├── coms-gfp.sh ├── coms-hopper-ablate-conservatism.sh ├── coms-hopper-demo-conservatism.sh ├── coms-hopper.sh ├── coms-nas.sh ├── coms-superconductor.sh ├── coms-tf-bind-10.sh ├── coms-tf-bind-8.sh ├── coms-utr-ablate-conservatism.sh ├── coms-utr-demo-conservatism.sh ├── coms-utr.sh ├── launch_chembl.sh ├── launch_demo.sh ├── launch_nas.sh └── launch_tf_bind_10.sh └── setup.py /.gitignore: -------------------------------------------------------------------------------- 1 | *.tfevents** 2 | *.egg-info** 3 | *idea** 4 | *data** 5 | *src** 6 | *.out** 7 | *.pkl** 8 | *.txt** 9 | *.json** 10 | *.csv** 11 | *.npy** 12 | *.pyc** -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2020 Brandon Trabucco 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Design-Baselines 2 | 3 | Design-Baselines is a set of **baseline algorithms** for solving automatic design problems that involve choosing an input that maximizes a black-box function. This type of optimization is used across scientific and engineering disciplines in ways such as designing proteins and DNA sequences with particular functions, chemical formulas and molecule substructures, the morphology and controllers of robots, and many more applications. 4 | 5 | These applications have significant potential to accelerate research in biochemistry, chemical engineering, materials science, robotics and many other disciplines. We hope this set of baselines serves as a robust platform to drive these applications and create widespread excitement for model-based optimization. 6 | 7 | **COMs Website**: [link](https://sites.google.com/berkeley.edu/coms/home?authuser=0) | **COMs Paper**: [arXiv](https://arxiv.org/abs/2107.06882) 8 | 9 | If these applications interest you, consider using our benchmark: [design-bench](https://github.com/brandontrabucco/design-bench), which you may install to python and automatically download benchmark data using `pip install design-bench==2.0.12`. 10 | 11 | ## Offline Model-Based Optimization 12 | 13 | ![Offline Model-Based Optimization](https://storage.googleapis.com/design-bench/mbo.png) 14 | 15 | The goal of model-based optimization is to find an input **x** that maximizes an unknown black-box function **f**. This function is frequently difficulty or costly to evaluate---such as requiring wet-lab experiments in the case of protein design. In these cases, **f** is described by a set of function evaluations: D = {(x_0, y_0), (x_1, y_1), ... (x_n, y_n)}, and optimization is performed without querying **f** on new data points. 16 | 17 | ## Installation 18 | 19 | Design-Baselines can be downloaded from github and installed using anaconda. 20 | 21 | ```bash 22 | git clone git@github.com:brandontrabucco/design-baselines.git 23 | conda create -f design-baselines/environment.yml 24 | ``` 25 | 26 | ## Performance Of Baselines 27 | 28 | We benchmark a set of 9 methods for solving offline model-based optimization problems. Performance is reported in normalized form, where the 100th percentile score of 128 candidate designs is evaluated and normalized such that a 1.0 corresponds to performance equivalent to the best performing design in the *full unobserved* dataset assoctated with each model-based optimization task. A 0.0 corresponds to performance equivalent to the worst performing design in the *full unobserved* dataset. In circumstances where an exact oracle is not available, this *full unobserved* dataset is used for training the approximate oracle that is used for evaluation of candidate designs proposed by each method. The symbol ± indicates the empirical standard deviation of reported performance across 8 trials. 29 | 30 | ### Performance On Continuous Tasks 31 | 32 | Method \ Task | Superconductor | Ant Morphology | D'Kitty Morphology | Hopper Controller 33 | ----------------------------- | -------------- | -------------- | ------------------ | ----------------- 34 | D (best) | 0.399 | 0.565 | 0.884 | 1.0 35 | Auto. CbAS | 0.421 ± 0.045 | 0.884 ± 0.046 | 0.906 ± 0.006 | 0.137 ± 0.005 36 | CbAS | 0.503 ± 0.069 | 0.879 ± 0.032 | 0.892 ± 0.008 | 0.141 ± 0.012 37 | BO-qEI | 0.402 ± 0.034 | 0.820 ± 0.000 | 0.896 ± 0.000 | 0.550 ± 0.118 38 | CMA-ES | 0.465 ± 0.024 | 1.219 ± 0.738 | 0.724 ± 0.001 | 0.604 ± 0.215 39 | Grad. | 0.518 ± 0.024 | 0.291 ± 0.023 | 0.874 ± 0.022 | 1.035 ± 0.482 40 | Grad. Min | 0.506 ± 0.009 | 0.478 ± 0.064 | 0.889 ± 0.011 | 1.391 ± 0.589 41 | Grad. Mean | 0.499 ± 0.017 | 0.444 ± 0.081 | 0.892 ± 0.011 | 1.586 ± 0.454 42 | MINs | 0.469 ± 0.023 | 0.916 ± 0.036 | 0.945 ± 0.012 | 0.424 ± 0.166 43 | REINFORCE | 0.481 ± 0.013 | 0.263 ± 0.032 | 0.562 ± 0.196 | -0.020 ± 0.067 44 | **COMs (Ours)** | 0.439 ± 0.033 | 0.944 ± 0.016 | 0.949 ± 0.015 | 2.056 ± 0.314 45 | 46 | ### Performance On Discrete Tasks 47 | 48 | Method \ Task | GFP | TF Bind 8 | UTR 49 | ----------------------------- | -------------- | -------------- | -------------- 50 | D (best) | 0.789 | 0.439 | 0.593 51 | Auto. CbAS | 0.865 ± 0.000 | 0.910 ± 0.044 | 0.691 ± 0.012 52 | CbAS | 0.865 ± 0.000 | 0.927 ± 0.051 | 0.694 ± 0.010 53 | BO-qEI | 0.254 ± 0.352 | 0.798 ± 0.083 | 0.684 ± 0.000 54 | CMA-ES | 0.054 ± 0.002 | 0.953 ± 0.022 | 0.707 ± 0.014 55 | Grad. | 0.864 ± 0.001 | 0.977 ± 0.025 | 0.695 ± 0.013 56 | Grad. Min | 0.864 ± 0.000 | 0.984 ± 0.012 | 0.696 ± 0.009 57 | Grad. Mean | 0.864 ± 0.000 | 0.986 ± 0.012 | 0.693 ± 0.010 58 | MINs | 0.865 ± 0.001 | 0.905 ± 0.052 | 0.697 ± 0.010 59 | REINFORCE | 0.865 ± 0.000 | 0.948 ± 0.028 | 0.688 ± 0.010 60 | **COMs (Ours)** | 0.864 ± 0.000 | 0.945 ± 0.033 | 0.699 ± 0.011 61 | 62 | ## Reproducing Baseline Performance 63 | 64 | To reproduce the performance of baseline algorithms reported in our work, you may then run the following series of commands in a bash terminal using the command-line interface exposed in design-baselines. Also, please ensure that the conda environment `design-baselines` is activated in the bash session that you run these commands from in order to access the `design-baselines` command-line interface. 65 | 66 | ```bash 67 | # set up machine parameters 68 | NUM_CPUS=32 69 | NUM_GPUS=8 70 | 71 | for TASK_NAME in \ 72 | superconductor \ 73 | ant \ 74 | dkitty \ 75 | hopper \ 76 | gfp \ 77 | tf-bind-8 \ 78 | utr; do 79 | 80 | for ALGORITHM_NAME in \ 81 | autofocused-cbas \ 82 | cbas \ 83 | bo-qei \ 84 | cma-es \ 85 | gradient-ascent \ 86 | gradient-ascent-min-ensemble \ 87 | gradient-ascent-mean-ensemble \ 88 | mins \ 89 | reinforce; do 90 | 91 | # launch several model-based optimization algorithms using the command line interface 92 | # for example: 93 | # (design-baselines) name@computer:~/$ cbas gfp \ 94 | # --local-dir ~/db-results/cbas-gfp \ 95 | # --cpus 32 \ 96 | # --gpus 8 \ 97 | # --num-parallel 8 \ 98 | # --num-samples 8 99 | $ALGORITHM_NAME $TASK_NAME \ 100 | --local-dir ~/db-results/$ALGORITHM_NAME-$TASK_NAME \ 101 | --cpus $NUM_CPUS \ 102 | --gpus $NUM_GPUS \ 103 | --num-parallel 8 \ 104 | --num-samples 8 105 | 106 | done 107 | 108 | done 109 | 110 | # generate the main performance table of the paper 111 | design-baselines make-table --dir ~/db-results/ --percentile 100th 112 | 113 | # generate the performance tables in the appendix 114 | design-baselines make-table --dir ~/db-results/ --percentile 50th 115 | design-baselines make-table --dir ~/db-results/ --percentile 100th --no-normalize 116 | ``` 117 | 118 | These commands will run several model-based optimization algorithms (such as [CbAS](http://proceedings.mlr.press/v97/brookes19a.html)) contained in design-baselines on all tasks released with the design-bench benchmark, and will then generate three performance tables from those results, and print a latex rendition of these performance tables to stdout. 119 | 120 | ## Running COMs 121 | 122 | You may run COMs using the `design-baselines` command line interface in a bash session where the `design-baselines` anaconda environments is activated and the `design-baselines` pip package is installed. Below is an example command that will run COMs on the task `HopperController-Exact-v0` from [design-bench](https://github.com/brandontrabucco/design-bench). 123 | 124 | ```bash 125 | coms --logging-dir ./coms-hopper \ 126 | --not-fast \ 127 | --task HopperController-Exact-v0 \ 128 | --no-task-relabel \ 129 | --normalize-ys \ 130 | --normalize-xs \ 131 | --particle-lr 0.05 \ 132 | --particle-train-gradient-steps 50 \ 133 | --particle-evaluate-gradient-steps 50 \ 134 | --particle-entropy-coefficient 0.0 \ 135 | --forward-model-activations relu \ 136 | --forward-model-activations relu \ 137 | --forward-model-hidden-size 2048 \ 138 | --forward-model-lr 0.0003 \ 139 | --forward-model-alpha 0.1 \ 140 | --forward-model-alpha-lr 0.01 \ 141 | --forward-model-overestimation-limit 0.5 \ 142 | --forward-model-noise-std 0.0 \ 143 | --forward-model-batch-size 128 \ 144 | --forward-model-val-size 500 \ 145 | --forward-model-epochs 50 \ 146 | --evaluation-samples 128 147 | ``` 148 | -------------------------------------------------------------------------------- /create_slurm_scripts.py: -------------------------------------------------------------------------------- 1 | import argparse as ap 2 | import os 3 | 4 | 5 | if __name__ == "__main__": 6 | 7 | parser = ap.ArgumentParser("Generate MBO Scripts") 8 | 9 | parser.add_argument("--image-dir", type=str, 10 | default='/global/scratch/btrabucco/mbo.img') 11 | parser.add_argument("--scripts-dir", type=str, 12 | default='/global/scratch/btrabucco/scripts') 13 | parser.add_argument("--results-dir", type=str, 14 | default='/global/scratch/btrabucco/mbo-results') 15 | 16 | parser.add_argument("--slurm-account", type=str, default='co_rail') 17 | parser.add_argument("--slurm-partition", type=str, default='savio3_gpu') 18 | parser.add_argument("--slurm-qos", type=str, default='rail_gpu3_normal') 19 | parser.add_argument("--slurm-hours", type=int, default=24) 20 | parser.add_argument("--slurm-memory", type=int, default=80) 21 | 22 | parser.add_argument("--singularity-args", type=str, default="") 23 | 24 | parser.add_argument("--cpus", type=int, default=8) 25 | parser.add_argument("--gpus", type=int, default=4) 26 | parser.add_argument("--num-parallel", type=int, default=8) 27 | parser.add_argument("--num-samples", type=int, default=8) 28 | 29 | args = parser.parse_args() 30 | 31 | os.makedirs(args.scripts_dir, exist_ok=True) 32 | os.makedirs(args.results_dir, exist_ok=True) 33 | 34 | for baseline in ["autofocused-cbas", 35 | "bo-qei", 36 | "cbas", 37 | "cma-es", 38 | "gradient-ascent", 39 | "gradient-ascent-min-ensemble", 40 | "gradient-ascent-mean-ensemble", 41 | "mins", 42 | "reinforce"]: 43 | 44 | for task in ["tf-bind-10", 45 | "nas"]: 46 | 47 | run_script = f"""#!/bin/bash 48 | . /packages/anaconda3/etc/profile.d/conda.sh 49 | 50 | export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:$HOME/.mujoco/mujoco200_linux/bin 51 | export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:$HOME/.mujoco/mujoco200/bin 52 | export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:$HOME/.mujoco/mjpro150/bin 53 | export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:$HOME/.mujoco/mjpro131/bin 54 | 55 | conda activate design-baselines 56 | 57 | {baseline} {task} \ 58 | --local-dir {args.results_dir}/{baseline}-{task} \ 59 | --cpus {args.cpus} \ 60 | --gpus {args.gpus} \ 61 | --num-parallel {args.num_parallel} \ 62 | --num-samples {args.num_samples}""" 63 | 64 | with open(f"{args.scripts_dir}/" 65 | f"run_{baseline}_{task}.sh", "w") as f: 66 | f.write(run_script) 67 | 68 | launch_script = f"""#!/bin/bash 69 | #SBATCH --job-name={baseline}-{task} 70 | #SBATCH --account={args.slurm_account} 71 | #SBATCH --time={args.slurm_hours}:00:00 72 | #SBATCH --partition={args.slurm_partition} 73 | #SBATCH --qos={args.slurm_qos} 74 | #SBATCH --cpus-per-task={args.cpus} 75 | #SBATCH --mem={args.slurm_memory}G 76 | #SBATCH --gres=gpu:TITAN:{args.gpus} 77 | 78 | singularity exec --nv -w {args.singularity_args} \ 79 | {args.image_dir} \ 80 | /bin/bash \ 81 | {args.scripts_dir}/run_{baseline}_{task}.sh""" 82 | 83 | with open(f"{args.scripts_dir}/" 84 | f"launch_{baseline}_{task}.sh", "w") as f: 85 | f.write(launch_script) 86 | 87 | print(f"sbatch {args.scripts_dir}/launch_{baseline}_{task}.sh") 88 | -------------------------------------------------------------------------------- /design_baselines/__init__.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | for gpu in tf.config.experimental.list_physical_devices('GPU'): 3 | tf.config.experimental.set_memory_growth(gpu, True) 4 | -------------------------------------------------------------------------------- /design_baselines/autofocused_cbas/__init__.py: -------------------------------------------------------------------------------- 1 | from design_baselines.data import StaticGraphTask, build_pipeline 2 | from design_baselines.logger import Logger 3 | from design_baselines.autofocused_cbas.trainers import Ensemble 4 | from design_baselines.autofocused_cbas.trainers import WeightedVAE 5 | from design_baselines.autofocused_cbas.trainers import CBAS 6 | from design_baselines.autofocused_cbas.nets import ForwardModel 7 | from design_baselines.autofocused_cbas.nets import Encoder 8 | from design_baselines.autofocused_cbas.nets import DiscreteDecoder 9 | from design_baselines.autofocused_cbas.nets import ContinuousDecoder 10 | import tensorflow as tf 11 | import numpy as np 12 | import os 13 | 14 | 15 | def autofocused_cbas(config): 16 | """Optimize a design problem score using the algorithm CBAS 17 | otherwise known as Conditioning by Adaptive Sampling 18 | 19 | Args: 20 | 21 | config: dict 22 | a dictionary of hyper parameters such as the learning rate 23 | """ 24 | 25 | logger = Logger(config['logging_dir']) 26 | task = StaticGraphTask(config['task'], **config['task_kwargs']) 27 | if task.is_discrete: 28 | task.map_to_integers() 29 | 30 | if config['normalize_ys']: 31 | task.map_normalize_y() 32 | if config['normalize_xs']: 33 | task.map_normalize_x() 34 | 35 | x = task.x 36 | y = task.y 37 | 38 | # create the training task and logger 39 | train_data, val_data = build_pipeline( 40 | x=x, y=y, w=np.ones_like(y), 41 | val_size=config['val_size'], 42 | batch_size=config['ensemble_batch_size'], 43 | bootstraps=config['bootstraps']) 44 | 45 | # make several keras neural networks with two hidden layers 46 | forward_models = [ForwardModel( 47 | task, 48 | embedding_size=config['embedding_size'], 49 | hidden_size=config['hidden_size'], 50 | num_layers=config['num_layers'], 51 | initial_max_std=config['initial_max_std'], 52 | initial_min_std=config['initial_min_std']) 53 | for b in range(config['bootstraps'])] 54 | 55 | # create a trainer for a forward model with a conservative objective 56 | ensemble = Ensemble( 57 | forward_models, 58 | forward_model_optim=tf.keras.optimizers.Adam, 59 | forward_model_lr=config['ensemble_lr']) 60 | 61 | # train the model for an additional number of epochs 62 | ensemble.launch(train_data, 63 | val_data, 64 | logger, 65 | config['ensemble_epochs']) 66 | 67 | # determine which arcitecture for the decoder to use 68 | decoder = DiscreteDecoder \ 69 | if task.is_discrete else ContinuousDecoder 70 | 71 | # build the encoder and decoder distribution and the p model 72 | p_encoder = Encoder(task, config['latent_size'], 73 | embedding_size=config['embedding_size'], 74 | hidden_size=config['hidden_size'], 75 | num_layers=config['num_layers'], 76 | initial_max_std=config['initial_max_std'], 77 | initial_min_std=config['initial_min_std']) 78 | p_decoder = decoder(task, config['latent_size'], 79 | hidden_size=config['hidden_size'], 80 | num_layers=config['num_layers'], 81 | initial_max_std=config['initial_max_std'], 82 | initial_min_std=config['initial_min_std']) 83 | p_vae = WeightedVAE(p_encoder, p_decoder, 84 | vae_optim=tf.keras.optimizers.Adam, 85 | vae_lr=config['vae_lr'], 86 | vae_beta=config['vae_beta']) 87 | 88 | # build a weighted data set 89 | train_data, val_data = build_pipeline( 90 | x=x, y=y, w=np.ones_like(task.y), 91 | batch_size=config['vae_batch_size'], 92 | val_size=config['val_size']) 93 | 94 | # train the initial vae fit to the original data distribution 95 | p_vae.launch(train_data, 96 | val_data, 97 | logger, 98 | config['offline_epochs']) 99 | 100 | # build the encoder and decoder distribution and the p model 101 | q_encoder = Encoder(task, config['latent_size'], 102 | embedding_size=config['embedding_size'], 103 | hidden_size=config['hidden_size'], 104 | num_layers=config['num_layers'], 105 | initial_max_std=config['initial_max_std'], 106 | initial_min_std=config['initial_min_std']) 107 | q_decoder = decoder(task, config['latent_size'], 108 | hidden_size=config['hidden_size'], 109 | num_layers=config['num_layers'], 110 | initial_max_std=config['initial_max_std'], 111 | initial_min_std=config['initial_min_std']) 112 | q_vae = WeightedVAE(q_encoder, q_decoder, 113 | vae_optim=tf.keras.optimizers.Adam, 114 | vae_lr=config['vae_lr'], 115 | vae_beta=config['vae_beta']) 116 | 117 | # create the cbas importance weight generator 118 | cbas = CBAS(ensemble, 119 | p_vae, 120 | q_vae, 121 | latent_size=config['latent_size']) 122 | 123 | # train and validate the q_vae using online samples 124 | q_encoder.set_weights(p_encoder.get_weights()) 125 | q_decoder.set_weights(p_decoder.get_weights()) 126 | for i in range(config['iterations']): 127 | 128 | # generate an importance weighted dataset 129 | x_t, y_t, w = cbas.generate_data( 130 | config['online_batches'], 131 | config['vae_batch_size'], 132 | config['percentile']) 133 | 134 | # build a weighted data set 135 | train_data, val_data = build_pipeline( 136 | x=x_t.numpy(), 137 | y=y_t.numpy(), 138 | w=w.numpy(), 139 | batch_size=config['vae_batch_size'], 140 | val_size=config['val_size']) 141 | 142 | # train a vae fit using weighted maximum likelihood 143 | start_epoch = config['online_epochs'] * i + \ 144 | config['offline_epochs'] 145 | q_vae.launch(train_data, 146 | val_data, 147 | logger, 148 | config['online_epochs'], 149 | start_epoch=start_epoch) 150 | 151 | # autofocus the forward model using importance weights 152 | v = cbas.autofocus_weights( 153 | x, batch_size=config['ensemble_batch_size']) 154 | train_data, val_data = build_pipeline( 155 | x=x, y=y, w=v.numpy(), 156 | bootstraps=config['bootstraps'], 157 | batch_size=config['ensemble_batch_size'], 158 | val_size=config['val_size']) 159 | 160 | # train a vae fit using weighted maximum likelihood 161 | start_epoch = config['autofocus_epochs'] * i + \ 162 | config['ensemble_epochs'] 163 | ensemble.launch(train_data, 164 | val_data, 165 | logger, 166 | config['autofocus_epochs'], 167 | start_epoch=start_epoch) 168 | 169 | # sample designs from the prior 170 | z = tf.random.normal([config['solver_samples'], config['latent_size']]) 171 | q_dx = q_decoder.get_distribution(z, training=False) 172 | x_t = q_dx.sample() 173 | np.save(os.path.join(config["logging_dir"], 174 | f"solution.npy"), x_t.numpy()) 175 | if config["do_evaluation"]: 176 | score = task.predict(x_t) 177 | if task.is_normalized_y: 178 | score = task.denormalize_y(score) 179 | logger.record("score", 180 | score, 181 | config['iterations'], 182 | percentile=True) 183 | -------------------------------------------------------------------------------- /design_baselines/bo_qei/nets.py: -------------------------------------------------------------------------------- 1 | from tensorflow_probability import distributions as tfpd 2 | import tensorflow.keras.layers as tfkl 3 | import tensorflow.keras as keras 4 | import tensorflow as tf 5 | import numpy as np 6 | 7 | 8 | class SequentialVAE(tf.keras.Model): 9 | 10 | def __init__(self, task, hidden_size=64, latent_size=256, 11 | activation='relu', kernel_size=3, num_blocks=4): 12 | 13 | super(SequentialVAE, self).__init__() 14 | 15 | input_shape = task.input_shape 16 | shape_before_flat = [input_shape[0] // ( 17 | 2 ** (num_blocks - 1)), hidden_size] 18 | 19 | """ 20 | 21 | DEFINE AN ENCODER MODEL THAT DOWNSAMPLES 22 | 23 | """ 24 | 25 | # the input layer of a keras model 26 | x = input_layer = keras.Input(shape=input_shape) 27 | 28 | # build a model with an input layer and optional embedding 29 | x = tfkl.Embedding(task.num_classes, hidden_size)(x) 30 | 31 | # the exponent of a positional embedding 32 | inverse_frequency = 1.0 / (10000.0 ** (tf.range( 33 | 0.0, hidden_size, 2.0) / hidden_size))[tf.newaxis] 34 | 35 | # calculate a positional embedding to break symmetry 36 | pos = tf.range(0.0, tf.shape(x)[1], 1.0)[:, tf.newaxis] 37 | positional_embedding = tf.concat([ 38 | tf.math.sin(pos * inverse_frequency), 39 | tf.math.cos(pos * inverse_frequency)], axis=1)[tf.newaxis] 40 | 41 | # add the positional encoding 42 | x = tfkl.Add()([x, positional_embedding]) 43 | x = tfkl.LayerNormalization()(x) 44 | 45 | # add several residual blocks to the model 46 | for i in range(num_blocks): 47 | 48 | if i > 0: 49 | # downsample the input sequence by 2 50 | x = tf.keras.layers.AveragePooling1D(pool_size=2, 51 | padding='same')(x) 52 | 53 | # first convolution layer in a residual block 54 | h = tfkl.Conv1D(hidden_size, kernel_size, 55 | padding='same', activation=None)(x) 56 | h = tfkl.LayerNormalization()(h) 57 | h = tfkl.Activation(activation)(h) 58 | 59 | # second convolution layer in a residual block 60 | h = tfkl.Conv1D(hidden_size, kernel_size, 61 | padding='same', activation=None)(h) 62 | h = tfkl.LayerNormalization()(h) 63 | h = tfkl.Activation(activation)(h) 64 | 65 | # add a residual connection to the model 66 | x = tfkl.Add()([x, h]) 67 | 68 | # flatten the result and predict the params of a gaussian 69 | flattened_x = tfkl.Flatten()(x) 70 | latent_mean = tfkl.Dense(latent_size)(flattened_x) 71 | latent_standard_dev = tfkl.Dense( 72 | latent_size, activation=tf.exp)(flattened_x) 73 | 74 | # save the encoder as a keras model 75 | self.encoder_cnn = keras.Model( 76 | inputs=input_layer, 77 | outputs=[latent_mean, latent_standard_dev]) 78 | 79 | """ 80 | 81 | DEFINE A DECODER THAT UPSAMPLES 82 | 83 | """ 84 | 85 | # the input layer of a keras model 86 | x = input_layer = keras.Input(shape=[latent_size]) 87 | x = tfkl.Dense(np.prod(shape_before_flat))(x) 88 | x = tfkl.Reshape(shape_before_flat)(x) 89 | 90 | # add several residual blocks to the model 91 | for i in reversed(range(num_blocks)): 92 | 93 | if i > 0: 94 | # up-sample the sequence and handle 95 | x = tf.pad(tf.repeat(x, 2, axis=1), [[0, 0], [ 96 | 0, 97 | (input_shape[0] // (2 ** (i - 1))) % 2 98 | ], [0, 0]], mode="SYMMETRIC") 99 | 100 | # the exponent of a positional embedding 101 | inverse_frequency = 1.0 / (10000.0 ** (tf.range( 102 | 0.0, hidden_size, 2.0) / hidden_size))[tf.newaxis] 103 | 104 | # calculate a positional embedding to break symmetry 105 | pos = tf.range(0.0, tf.shape(x)[1], 1.0)[:, tf.newaxis] 106 | positional_embedding = tf.concat([ 107 | tf.math.sin(pos * inverse_frequency), 108 | tf.math.cos(pos * inverse_frequency)], axis=1)[tf.newaxis] 109 | 110 | # add the positional encoding 111 | h = tfkl.Add()([x, positional_embedding]) 112 | h = tfkl.LayerNormalization()(h) 113 | 114 | # first convolution layer in a residual block 115 | h = tfkl.Conv1D(hidden_size, kernel_size, 116 | padding='same', activation=None)(h) 117 | h = tfkl.LayerNormalization()(h) 118 | h = tfkl.Activation(activation)(h) 119 | 120 | # second convolution layer in a residual block 121 | h = tfkl.Conv1D(hidden_size, kernel_size, 122 | padding='same', activation=None)(h) 123 | h = tfkl.LayerNormalization()(h) 124 | h = tfkl.Activation(activation)(h) 125 | 126 | # add a residual connection to the model 127 | x = tfkl.Add()([x, h]) 128 | 129 | # flatten the result and predict the params of a gaussian 130 | logits = tfkl.Dense(task.num_classes)(x) 131 | 132 | # save the encoder as a keras model 133 | self.decoder_cnn = keras.Model( 134 | inputs=input_layer, outputs=logits) 135 | 136 | def encode(self, x_batch, training=False): 137 | mean, standard_dev = self.encoder_cnn(x_batch, training=training) 138 | return tfpd.MultivariateNormalDiag(loc=mean, scale_diag=standard_dev) 139 | 140 | def decode(self, z, training=False): 141 | logits = self.decoder_cnn(z, training=training) 142 | return tfpd.Categorical(logits=logits) 143 | 144 | def generate(self, z, training=False): 145 | logits = self.decoder_cnn(z, training=training) 146 | return tf.argmax(logits, axis=2, output_type=tf.int32) 147 | 148 | 149 | class ForwardModel(tf.keras.Sequential): 150 | """A Fully Connected Network with 2 trainable layers""" 151 | 152 | distribution = tfpd.Normal 153 | 154 | def __init__(self, input_shape, hidden_size=50, 155 | num_layers=1, initial_max_std=1.5, initial_min_std=0.5): 156 | """Create a fully connected architecture using keras that can process 157 | designs and predict a gaussian distribution over scores 158 | 159 | Args: 160 | 161 | task: StaticGraphTask 162 | a model-based optimization task 163 | embedding_size: int 164 | the size of the embedding matrix for discrete tasks 165 | hidden_size: int 166 | the global hidden size of the neural network 167 | num_layers: int 168 | the number of hidden layers 169 | initial_max_std: float 170 | the starting upper bound of the standard deviation 171 | initial_min_std: float 172 | the starting lower bound of the standard deviation 173 | 174 | """ 175 | 176 | self.max_logstd = tf.Variable(tf.fill([1, 1], np.log( 177 | initial_max_std).astype(np.float32)), trainable=True) 178 | self.min_logstd = tf.Variable(tf.fill([1, 1], np.log( 179 | initial_min_std).astype(np.float32)), trainable=True) 180 | 181 | layers = [] 182 | layers.append(tfkl.Flatten(input_shape=input_shape) 183 | if len(layers) == 0 else tfkl.Flatten()) 184 | for i in range(num_layers): 185 | layers.extend([tfkl.Dense(hidden_size), tfkl.LeakyReLU()]) 186 | 187 | layers.append(tfkl.Dense(2)) 188 | super(ForwardModel, self).__init__(layers) 189 | 190 | def get_params(self, inputs, **kwargs): 191 | """Return a dictionary of parameters for a particular distribution 192 | family such as the mean and variance of a gaussian 193 | 194 | Args: 195 | 196 | inputs: tf.Tensor 197 | a batch of training inputs shaped like [batch_size, channels] 198 | 199 | Returns: 200 | 201 | parameters: dict 202 | a dictionary that contains 'loc' and 'scale_diag' keys 203 | """ 204 | 205 | prediction = super(ForwardModel, self).__call__(inputs, **kwargs) 206 | mean, logstd = tf.split(prediction, 2, axis=-1) 207 | logstd = self.max_logstd - tf.nn.softplus(self.max_logstd - logstd) 208 | logstd = self.min_logstd + tf.nn.softplus(logstd - self.min_logstd) 209 | return {"loc": mean, "scale": tf.math.exp(logstd)} 210 | 211 | def get_distribution(self, inputs, **kwargs): 212 | """Return a distribution over the outputs of this model, for example 213 | a Multivariate Gaussian Distribution 214 | 215 | Args: 216 | 217 | inputs: tf.Tensor 218 | a batch of training inputs shaped like [batch_size, channels] 219 | 220 | Returns: 221 | 222 | distribution: tfp.distribution.Distribution 223 | a tensorflow probability distribution over outputs of the model 224 | """ 225 | 226 | return self.distribution(**self.get_params(inputs, **kwargs)) 227 | -------------------------------------------------------------------------------- /design_baselines/cbas/__init__.py: -------------------------------------------------------------------------------- 1 | from design_baselines.data import StaticGraphTask, build_pipeline 2 | from design_baselines.logger import Logger 3 | from design_baselines.cbas.trainers import Ensemble 4 | from design_baselines.cbas.trainers import WeightedVAE 5 | from design_baselines.cbas.trainers import CBAS 6 | from design_baselines.cbas.nets import ForwardModel 7 | from design_baselines.cbas.nets import Encoder 8 | from design_baselines.cbas.nets import DiscreteDecoder 9 | from design_baselines.cbas.nets import ContinuousDecoder 10 | import tensorflow as tf 11 | import numpy as np 12 | import os 13 | 14 | 15 | def cbas(config): 16 | """Optimize a design problem score using the algorithm CBAS 17 | otherwise known as Conditioning by Adaptive Sampling 18 | 19 | Args: 20 | 21 | config: dict 22 | a dictionary of hyper parameters such as the learning rate 23 | """ 24 | 25 | logger = Logger(config['logging_dir']) 26 | task = StaticGraphTask(config['task'], **config['task_kwargs']) 27 | if task.is_discrete: 28 | task.map_to_integers() 29 | 30 | if config['normalize_ys']: 31 | task.map_normalize_y() 32 | if config['normalize_xs']: 33 | task.map_normalize_x() 34 | 35 | x = task.x 36 | y = task.y 37 | 38 | # create the training task and logger 39 | train_data, val_data = build_pipeline( 40 | x=x, y=y, w=np.ones_like(y), 41 | val_size=config['val_size'], 42 | batch_size=config['ensemble_batch_size'], 43 | bootstraps=config['bootstraps']) 44 | 45 | # make several keras neural networks with two hidden layers 46 | forward_models = [ForwardModel( 47 | task, 48 | embedding_size=config['embedding_size'], 49 | hidden_size=config['hidden_size'], 50 | num_layers=config['num_layers'], 51 | initial_max_std=config['initial_max_std'], 52 | initial_min_std=config['initial_min_std']) 53 | for b in range(config['bootstraps'])] 54 | 55 | # create a trainer for a forward model with a conservative objective 56 | ensemble = Ensemble( 57 | forward_models, 58 | forward_model_optim=tf.keras.optimizers.Adam, 59 | forward_model_lr=config['ensemble_lr']) 60 | 61 | # train the model for an additional number of epochs 62 | ensemble.launch(train_data, 63 | val_data, 64 | logger, 65 | config['ensemble_epochs']) 66 | 67 | # determine which arcitecture for the decoder to use 68 | decoder = DiscreteDecoder \ 69 | if task.is_discrete else ContinuousDecoder 70 | 71 | # build the encoder and decoder distribution and the p model 72 | p_encoder = Encoder(task, config['latent_size'], 73 | embedding_size=config['embedding_size'], 74 | hidden_size=config['hidden_size'], 75 | num_layers=config['num_layers'], 76 | initial_max_std=config['initial_max_std'], 77 | initial_min_std=config['initial_min_std']) 78 | p_decoder = decoder(task, config['latent_size'], 79 | hidden_size=config['hidden_size'], 80 | num_layers=config['num_layers'], 81 | initial_max_std=config['initial_max_std'], 82 | initial_min_std=config['initial_min_std']) 83 | p_vae = WeightedVAE(p_encoder, p_decoder, 84 | vae_optim=tf.keras.optimizers.Adam, 85 | vae_lr=config['vae_lr'], 86 | vae_beta=config['vae_beta']) 87 | 88 | # build a weighted data set 89 | train_data, val_data = build_pipeline( 90 | x=x, y=y, w=np.ones_like(task.y), 91 | batch_size=config['vae_batch_size'], 92 | val_size=config['val_size']) 93 | 94 | # train the initial vae fit to the original data distribution 95 | p_vae.launch(train_data, 96 | val_data, 97 | logger, 98 | config['offline_epochs']) 99 | 100 | # build the encoder and decoder distribution and the p model 101 | q_encoder = Encoder(task, config['latent_size'], 102 | embedding_size=config['embedding_size'], 103 | hidden_size=config['hidden_size'], 104 | num_layers=config['num_layers'], 105 | initial_max_std=config['initial_max_std'], 106 | initial_min_std=config['initial_min_std']) 107 | q_decoder = decoder(task, config['latent_size'], 108 | hidden_size=config['hidden_size'], 109 | num_layers=config['num_layers'], 110 | initial_max_std=config['initial_max_std'], 111 | initial_min_std=config['initial_min_std']) 112 | q_vae = WeightedVAE(q_encoder, q_decoder, 113 | vae_optim=tf.keras.optimizers.Adam, 114 | vae_lr=config['vae_lr'], 115 | vae_beta=config['vae_beta']) 116 | 117 | # create the cbas importance weight generator 118 | cbas = CBAS(ensemble, 119 | p_vae, 120 | q_vae, 121 | latent_size=config['latent_size']) 122 | 123 | # train and validate the q_vae using online samples 124 | q_encoder.set_weights(p_encoder.get_weights()) 125 | q_decoder.set_weights(p_decoder.get_weights()) 126 | for i in range(config['iterations']): 127 | 128 | # generate an importance weighted dataset 129 | x_t, y_t, w = cbas.generate_data( 130 | config['online_batches'], 131 | config['vae_batch_size'], 132 | config['percentile']) 133 | 134 | # build a weighted data set 135 | train_data, val_data = build_pipeline( 136 | x=x_t.numpy(), 137 | y=y_t.numpy(), 138 | w=w.numpy(), 139 | batch_size=config['vae_batch_size'], 140 | val_size=config['val_size']) 141 | 142 | # train a vae fit using weighted maximum likelihood 143 | start_epoch = config['online_epochs'] * i + \ 144 | config['offline_epochs'] 145 | q_vae.launch(train_data, 146 | val_data, 147 | logger, 148 | config['online_epochs'], 149 | start_epoch=start_epoch) 150 | 151 | # sample designs from the prior 152 | z = tf.random.normal([config['solver_samples'], config['latent_size']]) 153 | q_dx = q_decoder.get_distribution(z, training=False) 154 | x_t = q_dx.sample() 155 | np.save(os.path.join(config["logging_dir"], 156 | f"solution.npy"), x_t.numpy()) 157 | if config["do_evaluation"]: 158 | score = task.predict(x_t) 159 | if task.is_normalized_y: 160 | score = task.denormalize_y(score) 161 | logger.record("score", 162 | score, 163 | config['iterations'], 164 | percentile=True) 165 | -------------------------------------------------------------------------------- /design_baselines/cbas/distribution_experiments.py: -------------------------------------------------------------------------------- 1 | from ray import tune 2 | import click 3 | import ray 4 | import os 5 | 6 | 7 | @click.group() 8 | def cli(): 9 | """A group of experiments for training Conservative Score Models 10 | and reproducing our ICLR 2021 results. 11 | """ 12 | 13 | 14 | ############# 15 | 16 | 17 | @cli.command() 18 | @click.option('--local-dir', type=str, default='cbas-hopper') 19 | @click.option('--cpus', type=int, default=24) 20 | @click.option('--gpus', type=int, default=1) 21 | @click.option('--num-parallel', type=int, default=1) 22 | @click.option('--num-samples', type=int, default=1) 23 | def hopper(local_dir, cpus, gpus, num_parallel, num_samples): 24 | """Evaluate AutoFocusing on HopperController-Exact-v0 25 | """ 26 | 27 | # Final Version 28 | 29 | from design_baselines.cbas import cbas 30 | ray.init(num_cpus=cpus, 31 | num_gpus=gpus, 32 | include_dashboard=False, 33 | _temp_dir=os.path.expanduser('~/tmp')) 34 | tune.run(cbas, config={ 35 | "logging_dir": "data", 36 | "normalize_ys": True, 37 | "normalize_xs": True, 38 | "task": "HopperController-Exact-v0", 39 | "task_kwargs": {"relabel": False, "dataset_kwargs": { 40 | "max_samples": 1000, 41 | "distribution": tune.grid_search([ 42 | "uniform", 43 | "linear", 44 | "quadratic", 45 | "exponential", 46 | "circular" 47 | ]), 48 | "max_percentile": 100, 49 | "min_percentile": 0 50 | }}, 51 | "bootstraps": 5, 52 | "val_size": 200, 53 | "ensemble_batch_size": 100, 54 | "vae_batch_size": 100, 55 | "embedding_size": 256, 56 | "hidden_size": 256, 57 | "num_layers": 1, 58 | "initial_max_std": 0.2, 59 | "initial_min_std": 0.1, 60 | "ensemble_lr": 0.0003, 61 | "ensemble_epochs": 100, 62 | "latent_size": 32, 63 | "vae_lr": 0.0003, 64 | "vae_beta": 1.0, 65 | "offline_epochs": 200, 66 | "online_batches": 10, 67 | "online_epochs": 10, 68 | "iterations": 20, 69 | "percentile": 80.0, 70 | "solver_samples": 512}, 71 | num_samples=num_samples, 72 | local_dir=local_dir, 73 | resources_per_trial={'cpu': cpus // num_parallel, 74 | 'gpu': gpus / num_parallel - 0.01}) 75 | 76 | 77 | @cli.command() 78 | @click.option('--local-dir', type=str, default='cbas-superconductor') 79 | @click.option('--cpus', type=int, default=24) 80 | @click.option('--gpus', type=int, default=1) 81 | @click.option('--num-parallel', type=int, default=1) 82 | @click.option('--num-samples', type=int, default=1) 83 | @click.option('--oracle', type=str, default="RandomForest") 84 | def superconductor(local_dir, cpus, gpus, num_parallel, num_samples, oracle): 85 | """Evaluate AutoFocusing on Superconductor-RandomForest-v0 86 | """ 87 | 88 | # Final Version 89 | 90 | from design_baselines.cbas import cbas 91 | ray.init(num_cpus=cpus, 92 | num_gpus=gpus, 93 | include_dashboard=False, 94 | _temp_dir=os.path.expanduser('~/tmp')) 95 | tune.run(cbas, config={ 96 | "logging_dir": "data", 97 | "normalize_ys": True, 98 | "normalize_xs": True, 99 | "task": f"Superconductor-{oracle}-v0", 100 | "task_kwargs": {"relabel": False, "dataset_kwargs": { 101 | "max_samples": 5000, 102 | "distribution": tune.grid_search([ 103 | "uniform", 104 | "linear", 105 | "quadratic", 106 | "exponential", 107 | "circular" 108 | ]), 109 | "max_percentile": 100, 110 | "min_percentile": 0 111 | }}, 112 | "bootstraps": 5, 113 | "val_size": 200, 114 | "ensemble_batch_size": 100, 115 | "vae_batch_size": 100, 116 | "embedding_size": 256, 117 | "hidden_size": 256, 118 | "num_layers": 1, 119 | "initial_max_std": 0.2, 120 | "initial_min_std": 0.1, 121 | "ensemble_lr": 0.0003, 122 | "ensemble_epochs": 100, 123 | "latent_size": 32, 124 | "vae_lr": 0.0003, 125 | "vae_beta": 1.0, 126 | "offline_epochs": 200, 127 | "online_batches": 10, 128 | "online_epochs": 10, 129 | "iterations": 20, 130 | "percentile": 80.0, 131 | "solver_samples": 512}, 132 | num_samples=num_samples, 133 | local_dir=local_dir, 134 | resources_per_trial={'cpu': cpus // num_parallel, 135 | 'gpu': gpus / num_parallel - 0.01}) 136 | 137 | 138 | @cli.command() 139 | @click.option('--local-dir', type=str, default='cbas-gfp') 140 | @click.option('--cpus', type=int, default=24) 141 | @click.option('--gpus', type=int, default=1) 142 | @click.option('--num-parallel', type=int, default=1) 143 | @click.option('--num-samples', type=int, default=1) 144 | @click.option('--oracle', type=str, default="Transformer") 145 | def gfp(local_dir, cpus, gpus, num_parallel, num_samples, oracle): 146 | """Evaluate AutoFocusing on GFP-Transformer-v0 147 | """ 148 | 149 | # Final Version 150 | 151 | from design_baselines.cbas import cbas 152 | ray.init(num_cpus=cpus, 153 | num_gpus=gpus, 154 | include_dashboard=False, 155 | _temp_dir=os.path.expanduser('~/tmp')) 156 | tune.run(cbas, config={ 157 | "logging_dir": "data", 158 | "normalize_ys": True, 159 | "normalize_xs": False, 160 | "task": f"GFP-{oracle}-v0", 161 | "task_kwargs": {"relabel": False, "dataset_kwargs": { 162 | "max_samples": 5000, 163 | "distribution": tune.grid_search([ 164 | "uniform", 165 | "linear", 166 | "quadratic", 167 | "exponential", 168 | "circular" 169 | ]), 170 | "max_percentile": 100, 171 | "min_percentile": 0 172 | }}, 173 | "bootstraps": 5, 174 | "val_size": 200, 175 | "ensemble_batch_size": 100, 176 | "vae_batch_size": 100, 177 | "embedding_size": 256, 178 | "hidden_size": 256, 179 | "num_layers": 1, 180 | "initial_max_std": 0.2, 181 | "initial_min_std": 0.1, 182 | "ensemble_lr": 0.0003, 183 | "ensemble_epochs": 100, 184 | "latent_size": 32, 185 | "vae_lr": 0.0003, 186 | "vae_beta": 1.0, 187 | "offline_epochs": 200, 188 | "online_batches": 10, 189 | "online_epochs": 10, 190 | "iterations": 20, 191 | "percentile": 80.0, 192 | "solver_samples": 512}, 193 | num_samples=num_samples, 194 | local_dir=local_dir, 195 | resources_per_trial={'cpu': cpus // num_parallel, 196 | 'gpu': gpus / num_parallel - 0.01}) 197 | 198 | 199 | @cli.command() 200 | @click.option('--local-dir', type=str, default='cbas-utr') 201 | @click.option('--cpus', type=int, default=24) 202 | @click.option('--gpus', type=int, default=1) 203 | @click.option('--num-parallel', type=int, default=1) 204 | @click.option('--num-samples', type=int, default=1) 205 | def utr(local_dir, cpus, gpus, num_parallel, num_samples): 206 | """Evaluate AutoFocusing on UTR-ResNet-v0 207 | """ 208 | 209 | # Final Version 210 | 211 | from design_baselines.cbas import cbas 212 | ray.init(num_cpus=cpus, 213 | num_gpus=gpus, 214 | include_dashboard=False, 215 | _temp_dir=os.path.expanduser('~/tmp')) 216 | tune.run(cbas, config={ 217 | "logging_dir": "data", 218 | "normalize_ys": True, 219 | "normalize_xs": False, 220 | "task": "UTR-ResNet-v0", 221 | "task_kwargs": {"relabel": False, "dataset_kwargs": { 222 | "max_samples": 20000, 223 | "distribution": tune.grid_search([ 224 | "uniform", 225 | "linear", 226 | "quadratic", 227 | "exponential", 228 | "circular" 229 | ]), 230 | "max_percentile": 100, 231 | "min_percentile": 0 232 | }}, 233 | "bootstraps": 5, 234 | "val_size": 200, 235 | "ensemble_batch_size": 100, 236 | "vae_batch_size": 100, 237 | "embedding_size": 256, 238 | "hidden_size": 256, 239 | "num_layers": 1, 240 | "initial_max_std": 0.2, 241 | "initial_min_std": 0.1, 242 | "ensemble_lr": 0.0003, 243 | "ensemble_epochs": 100, 244 | "latent_size": 32, 245 | "vae_lr": 0.0003, 246 | "vae_beta": 1.0, 247 | "offline_epochs": 200, 248 | "online_batches": 10, 249 | "online_epochs": 10, 250 | "iterations": 20, 251 | "percentile": 80.0, 252 | "solver_samples": 512}, 253 | num_samples=num_samples, 254 | local_dir=local_dir, 255 | resources_per_trial={'cpu': cpus // num_parallel, 256 | 'gpu': gpus / num_parallel - 0.01}) 257 | 258 | -------------------------------------------------------------------------------- /design_baselines/cma_es/__init__.py: -------------------------------------------------------------------------------- 1 | from design_baselines.data import StaticGraphTask, build_pipeline 2 | from design_baselines.logger import Logger 3 | from design_baselines.cma_es.trainers import Ensemble, VAETrainer 4 | from design_baselines.cma_es.nets import ForwardModel, SequentialVAE 5 | from tensorflow_probability import distributions as tfpd 6 | import tensorflow as tf 7 | import tensorflow.keras as keras 8 | import os 9 | import math 10 | import numpy as np 11 | 12 | 13 | def cma_es(config): 14 | """Optimizes over designs x in an offline optimization problem 15 | using the CMA Evolution Strategy 16 | 17 | Args: 18 | 19 | config: dict 20 | a dictionary of hyper parameters such as the learning rate 21 | """ 22 | 23 | # create the training task and logger 24 | logger = Logger(config['logging_dir']) 25 | task = StaticGraphTask(config['task'], **config['task_kwargs']) 26 | 27 | if config['normalize_ys']: 28 | task.map_normalize_y() 29 | if task.is_discrete and not config["use_vae"]: 30 | task.map_to_logits() 31 | if config['normalize_xs']: 32 | task.map_normalize_x() 33 | 34 | x = task.x 35 | y = task.y 36 | 37 | if task.is_discrete and config["use_vae"]: 38 | 39 | vae_model = SequentialVAE( 40 | task, 41 | hidden_size=config['vae_hidden_size'], 42 | latent_size=config['vae_latent_size'], 43 | activation=config['vae_activation'], 44 | kernel_size=config['vae_kernel_size'], 45 | num_blocks=config['vae_num_blocks']) 46 | 47 | vae_trainer = VAETrainer(vae_model, 48 | vae_optim=tf.keras.optimizers.Adam, 49 | vae_lr=config['vae_lr'], 50 | beta=config['vae_beta']) 51 | 52 | # create the training task and logger 53 | train_data, val_data = build_pipeline( 54 | x=x, y=y, 55 | batch_size=config['vae_batch_size'], 56 | val_size=config['val_size']) 57 | 58 | # estimate the number of training steps per epoch 59 | vae_trainer.launch(train_data, val_data, 60 | logger, config['vae_epochs']) 61 | 62 | # map the x values to latent space 63 | x = vae_model.encoder_cnn.predict(x)[0] 64 | 65 | mean = np.mean(x, axis=0, keepdims=True) 66 | standard_dev = np.std(x - mean, axis=0, keepdims=True) 67 | x = (x - mean) / standard_dev 68 | 69 | input_shape = x.shape[1:] 70 | input_size = np.prod(input_shape) 71 | 72 | # make several keras neural networks with two hidden layers 73 | forward_models = [ForwardModel( 74 | input_shape, 75 | hidden_size=config['hidden_size'], 76 | num_layers=config['num_layers'], 77 | initial_max_std=config['initial_max_std'], 78 | initial_min_std=config['initial_min_std']) 79 | for b in range(config['bootstraps'])] 80 | 81 | # create a trainer for a forward model with a conservative objective 82 | ensemble = Ensemble( 83 | forward_models, 84 | forward_model_optim=tf.keras.optimizers.Adam, 85 | forward_model_lr=config['ensemble_lr']) 86 | 87 | # create the training task and logger 88 | train_data, val_data = build_pipeline( 89 | x=x, y=y, bootstraps=config['bootstraps'], 90 | batch_size=config['ensemble_batch_size'], 91 | val_size=config['val_size']) 92 | 93 | # train the model for an additional number of epochs 94 | ensemble.launch(train_data, 95 | val_data, 96 | logger, 97 | config['ensemble_epochs']) 98 | 99 | # select the top 1 initial designs from the dataset 100 | indices = tf.math.top_k(y[:, 0], k=config['solver_samples'])[1] 101 | initial_x = tf.gather(x, indices, axis=0) 102 | x = initial_x 103 | 104 | # create a fitness function for optimizing the expected task score 105 | def fitness(input_x): 106 | input_x = tf.reshape(input_x, input_shape)[tf.newaxis] 107 | if config["optimize_ground_truth"]: 108 | if task.is_discrete and config["use_vae"]: 109 | input_x = tf.argmax(vae_model.decoder_cnn.predict( 110 | input_x * standard_dev + mean), axis=2, output_type=tf.int32) 111 | value = task.predict(input_x) 112 | else: 113 | value = ensemble.get_distribution(input_x).mean() 114 | return (-value[0].numpy()).tolist()[0] 115 | 116 | import cma 117 | result = [] 118 | for i in range(config['solver_samples']): 119 | xi = x[i].numpy().flatten().tolist() 120 | es = cma.CMAEvolutionStrategy(xi, config['cma_sigma']) 121 | step = 0 122 | while not es.stop() and step < config['cma_max_iterations']: 123 | solutions = es.ask() 124 | es.tell(solutions, [fitness(x) for x in solutions]) 125 | step += 1 126 | result.append( 127 | tf.reshape(es.result.xbest, input_shape)) 128 | print(f"CMA: {i + 1} / {config['solver_samples']}") 129 | 130 | # convert the solution found by CMA-ES to a tensor 131 | x = tf.stack(result, axis=0) 132 | solution = x 133 | 134 | if task.is_discrete and config["use_vae"]: 135 | solution = solution * standard_dev + mean 136 | logits = vae_model.decoder_cnn.predict(solution) 137 | solution = tf.argmax(logits, axis=2, output_type=tf.int32) 138 | 139 | # save the current solution to the disk 140 | np.save(os.path.join(config["logging_dir"], 141 | f"solution.npy"), solution.numpy()) 142 | if config["do_evaluation"]: 143 | 144 | # evaluate the found solution 145 | score = task.predict(solution) 146 | if task.is_normalized_y: 147 | score = task.denormalize_y(score) 148 | logger.record("score", score, 0, percentile=True) 149 | -------------------------------------------------------------------------------- /design_baselines/cma_es/distribution_experiments.py: -------------------------------------------------------------------------------- 1 | from ray import tune 2 | import click 3 | import ray 4 | import os 5 | 6 | 7 | @click.group() 8 | def cli(): 9 | """A group of experiments for training Conservative Score Models 10 | and reproducing our ICLR 2021 results. 11 | """ 12 | 13 | 14 | ############# 15 | 16 | 17 | @cli.command() 18 | @click.option('--local-dir', type=str, default='cma-es-hopper') 19 | @click.option('--cpus', type=int, default=24) 20 | @click.option('--gpus', type=int, default=1) 21 | @click.option('--num-parallel', type=int, default=1) 22 | @click.option('--num-samples', type=int, default=1) 23 | def hopper(local_dir, cpus, gpus, num_parallel, num_samples): 24 | """Evaluate CMA-ES on HopperController-Exact-v0 25 | """ 26 | 27 | # Final Version 28 | 29 | from design_baselines.cma_es import cma_es 30 | ray.init(num_cpus=cpus, 31 | num_gpus=gpus, 32 | include_dashboard=False, 33 | _temp_dir=os.path.expanduser('~/tmp')) 34 | tune.run(cma_es, config={ 35 | "logging_dir": "data", 36 | "normalize_ys": True, 37 | "normalize_xs": True, 38 | "task": "HopperController-Exact-v0", 39 | "task_kwargs": {"relabel": False, "dataset_kwargs": { 40 | "max_samples": 1000, 41 | "distribution": tune.grid_search([ 42 | "uniform", 43 | "linear", 44 | "quadratic", 45 | "exponential", 46 | "circular" 47 | ]), 48 | "max_percentile": 100, 49 | "min_percentile": 0 50 | }}, 51 | "bootstraps": 5, 52 | "val_size": 200, 53 | "optimize_ground_truth": False, 54 | "ensemble_batch_size": 100, 55 | "hidden_size": 256, 56 | "num_layers": 1, 57 | "initial_max_std": 0.2, 58 | "initial_min_std": 0.1, 59 | "ensemble_lr": 0.001, 60 | "ensemble_epochs": 100, 61 | "cma_max_iterations": 100, 62 | "cma_sigma": 0.5, 63 | "solver_samples": 512}, 64 | num_samples=num_samples, 65 | local_dir=local_dir, 66 | resources_per_trial={'cpu': cpus // num_parallel, 67 | 'gpu': gpus / num_parallel - 0.01}) 68 | 69 | 70 | @cli.command() 71 | @click.option('--local-dir', type=str, default='cma-es-superconductor') 72 | @click.option('--cpus', type=int, default=24) 73 | @click.option('--gpus', type=int, default=1) 74 | @click.option('--num-parallel', type=int, default=1) 75 | @click.option('--num-samples', type=int, default=1) 76 | @click.option('--oracle', type=str, default="RandomForest") 77 | def superconductor(local_dir, cpus, gpus, num_parallel, num_samples, oracle): 78 | """Evaluate CMA-ES on Superconductor-RandomForest-v0 79 | """ 80 | 81 | # Final Version 82 | 83 | from design_baselines.cma_es import cma_es 84 | ray.init(num_cpus=cpus, 85 | num_gpus=gpus, 86 | include_dashboard=False, 87 | _temp_dir=os.path.expanduser('~/tmp')) 88 | tune.run(cma_es, config={ 89 | "logging_dir": "data", 90 | "normalize_ys": True, 91 | "normalize_xs": True, 92 | "task": f"Superconductor-{oracle}-v0", 93 | "task_kwargs": {"relabel": False, "dataset_kwargs": { 94 | "max_samples": 5000, 95 | "distribution": tune.grid_search([ 96 | "uniform", 97 | "linear", 98 | "quadratic", 99 | "exponential", 100 | "circular" 101 | ]), 102 | "max_percentile": 100, 103 | "min_percentile": 0 104 | }}, 105 | "bootstraps": 5, 106 | "val_size": 200, 107 | "optimize_ground_truth": False, 108 | "ensemble_batch_size": 100, 109 | "hidden_size": 256, 110 | "num_layers": 1, 111 | "initial_max_std": 0.2, 112 | "initial_min_std": 0.1, 113 | "ensemble_lr": 0.001, 114 | "ensemble_epochs": 100, 115 | "cma_max_iterations": 100, 116 | "cma_sigma": 0.5, 117 | "solver_samples": 512}, 118 | num_samples=num_samples, 119 | local_dir=local_dir, 120 | resources_per_trial={'cpu': cpus // num_parallel, 121 | 'gpu': gpus / num_parallel - 0.01}) 122 | 123 | 124 | @cli.command() 125 | @click.option('--local-dir', type=str, default='cma-es-gfp') 126 | @click.option('--cpus', type=int, default=24) 127 | @click.option('--gpus', type=int, default=1) 128 | @click.option('--num-parallel', type=int, default=1) 129 | @click.option('--num-samples', type=int, default=1) 130 | @click.option('--oracle', type=str, default="Transformer") 131 | def gfp(local_dir, cpus, gpus, num_parallel, num_samples, oracle): 132 | """Evaluate CMA-ES on GFP-Transformer-v0 133 | """ 134 | 135 | # Final Version 136 | 137 | from design_baselines.cma_es import cma_es 138 | ray.init(num_cpus=cpus, 139 | num_gpus=gpus, 140 | include_dashboard=False, 141 | _temp_dir=os.path.expanduser('~/tmp')) 142 | tune.run(cma_es, config={ 143 | "logging_dir": "data", 144 | "normalize_ys": True, 145 | "normalize_xs": True, 146 | "task": f"GFP-{oracle}-v0", 147 | "task_kwargs": {"relabel": False, "dataset_kwargs": { 148 | "max_samples": 5000, 149 | "distribution": tune.grid_search([ 150 | "uniform", 151 | "linear", 152 | "quadratic", 153 | "exponential", 154 | "circular" 155 | ]), 156 | "max_percentile": 100, 157 | "min_percentile": 0 158 | }}, 159 | "bootstraps": 5, 160 | "val_size": 200, 161 | "optimize_ground_truth": False, 162 | "use_vae": False, 163 | "vae_beta": 0.01, 164 | "vae_epochs": 20, 165 | "vae_batch_size": 128, 166 | "vae_hidden_size": 64, 167 | "vae_latent_size": 256, 168 | "vae_activation": "relu", 169 | "vae_kernel_size": 3, 170 | "vae_num_blocks": 4, 171 | "vae_lr": 0.001, 172 | "ensemble_batch_size": 100, 173 | "hidden_size": 256, 174 | "num_layers": 1, 175 | "initial_max_std": 0.2, 176 | "initial_min_std": 0.1, 177 | "ensemble_lr": 0.001, 178 | "ensemble_epochs": 100, 179 | "cma_max_iterations": 100, 180 | "cma_sigma": 0.5, 181 | "solver_samples": 512}, 182 | num_samples=num_samples, 183 | local_dir=local_dir, 184 | resources_per_trial={'cpu': cpus // num_parallel, 185 | 'gpu': gpus / num_parallel - 0.01}) 186 | 187 | 188 | @cli.command() 189 | @click.option('--local-dir', type=str, default='cma-es-utr') 190 | @click.option('--cpus', type=int, default=24) 191 | @click.option('--gpus', type=int, default=1) 192 | @click.option('--num-parallel', type=int, default=1) 193 | @click.option('--num-samples', type=int, default=1) 194 | def utr(local_dir, cpus, gpus, num_parallel, num_samples): 195 | """Evaluate CMA-ES on UTR-ResNet-v0 196 | """ 197 | 198 | # Final Version 199 | 200 | from design_baselines.cma_es import cma_es 201 | ray.init(num_cpus=cpus, 202 | num_gpus=gpus, 203 | include_dashboard=False, 204 | _temp_dir=os.path.expanduser('~/tmp')) 205 | tune.run(cma_es, config={ 206 | "logging_dir": "data", 207 | "normalize_ys": True, 208 | "normalize_xs": True, 209 | "task": "UTR-ResNet-v0", 210 | "task_kwargs": {"relabel": False, "dataset_kwargs": { 211 | "max_samples": 20000, 212 | "distribution": tune.grid_search([ 213 | "uniform", 214 | "linear", 215 | "quadratic", 216 | "exponential", 217 | "circular" 218 | ]), 219 | "max_percentile": 100, 220 | "min_percentile": 0 221 | }}, 222 | "bootstraps": 5, 223 | "val_size": 200, 224 | "optimize_ground_truth": False, 225 | "use_vae": False, 226 | "vae_beta": 0.01, 227 | "vae_epochs": 20, 228 | "vae_batch_size": 128, 229 | "vae_hidden_size": 64, 230 | "vae_latent_size": 256, 231 | "vae_activation": "relu", 232 | "vae_kernel_size": 3, 233 | "vae_num_blocks": 4, 234 | "vae_lr": 0.001, 235 | "ensemble_batch_size": 100, 236 | "hidden_size": 256, 237 | "num_layers": 1, 238 | "initial_max_std": 0.2, 239 | "initial_min_std": 0.1, 240 | "ensemble_lr": 0.001, 241 | "ensemble_epochs": 100, 242 | "cma_max_iterations": 100, 243 | "cma_sigma": 0.5, 244 | "solver_samples": 512}, 245 | num_samples=num_samples, 246 | local_dir=local_dir, 247 | resources_per_trial={'cpu': cpus // num_parallel, 248 | 'gpu': gpus / num_parallel - 0.01}) 249 | 250 | -------------------------------------------------------------------------------- /design_baselines/cma_es/nets.py: -------------------------------------------------------------------------------- 1 | from tensorflow_probability import distributions as tfpd 2 | import tensorflow.keras.layers as tfkl 3 | import tensorflow.keras as keras 4 | import tensorflow as tf 5 | import numpy as np 6 | 7 | 8 | class SequentialVAE(tf.keras.Model): 9 | 10 | def __init__(self, task, hidden_size=64, latent_size=256, 11 | activation='relu', kernel_size=3, num_blocks=4): 12 | 13 | super(SequentialVAE, self).__init__() 14 | 15 | input_shape = task.input_shape 16 | shape_before_flat = [input_shape[0] // ( 17 | 2 ** (num_blocks - 1)), hidden_size] 18 | 19 | """ 20 | 21 | DEFINE AN ENCODER MODEL THAT DOWNSAMPLES 22 | 23 | """ 24 | 25 | # the input layer of a keras model 26 | x = input_layer = keras.Input(shape=input_shape) 27 | 28 | # build a model with an input layer and optional embedding 29 | x = tfkl.Embedding(task.num_classes, hidden_size)(x) 30 | 31 | # the exponent of a positional embedding 32 | inverse_frequency = 1.0 / (10000.0 ** (tf.range( 33 | 0.0, hidden_size, 2.0) / hidden_size))[tf.newaxis] 34 | 35 | # calculate a positional embedding to break symmetry 36 | pos = tf.range(0.0, tf.shape(x)[1], 1.0)[:, tf.newaxis] 37 | positional_embedding = tf.concat([ 38 | tf.math.sin(pos * inverse_frequency), 39 | tf.math.cos(pos * inverse_frequency)], axis=1)[tf.newaxis] 40 | 41 | # add the positional encoding 42 | x = tfkl.Add()([x, positional_embedding]) 43 | x = tfkl.LayerNormalization()(x) 44 | 45 | # add several residual blocks to the model 46 | for i in range(num_blocks): 47 | 48 | if i > 0: 49 | 50 | # downsample the input sequence by 2 51 | x = tf.keras.layers.AveragePooling1D(pool_size=2, 52 | padding='same')(x) 53 | 54 | # first convolution layer in a residual block 55 | h = tfkl.Conv1D(hidden_size, kernel_size, 56 | padding='same', activation=None)(x) 57 | h = tfkl.LayerNormalization()(h) 58 | h = tfkl.Activation(activation)(h) 59 | 60 | # second convolution layer in a residual block 61 | h = tfkl.Conv1D(hidden_size, kernel_size, 62 | padding='same', activation=None)(h) 63 | h = tfkl.LayerNormalization()(h) 64 | h = tfkl.Activation(activation)(h) 65 | 66 | # add a residual connection to the model 67 | x = tfkl.Add()([x, h]) 68 | 69 | # flatten the result and predict the params of a gaussian 70 | flattened_x = tfkl.Flatten()(x) 71 | latent_mean = tfkl.Dense(latent_size)(flattened_x) 72 | latent_standard_dev = tfkl.Dense( 73 | latent_size, activation=tf.exp)(flattened_x) 74 | 75 | # save the encoder as a keras model 76 | self.encoder_cnn = keras.Model( 77 | inputs=input_layer, 78 | outputs=[latent_mean, latent_standard_dev]) 79 | 80 | """ 81 | 82 | DEFINE A DECODER THAT UPSAMPLES 83 | 84 | """ 85 | 86 | # the input layer of a keras model 87 | x = input_layer = keras.Input(shape=[latent_size]) 88 | x = tfkl.Dense(np.prod(shape_before_flat))(x) 89 | x = tfkl.Reshape(shape_before_flat)(x) 90 | 91 | # add several residual blocks to the model 92 | for i in reversed(range(num_blocks)): 93 | 94 | if i > 0: 95 | 96 | # up-sample the sequence and handle 97 | x = tf.pad(tf.repeat(x, 2, axis=1), [[0, 0], [ 98 | 0, 99 | (input_shape[0] // (2 ** (i - 1))) % 2 100 | ], [0, 0]], mode="SYMMETRIC") 101 | 102 | # the exponent of a positional embedding 103 | inverse_frequency = 1.0 / (10000.0 ** (tf.range( 104 | 0.0, hidden_size, 2.0) / hidden_size))[tf.newaxis] 105 | 106 | # calculate a positional embedding to break symmetry 107 | pos = tf.range(0.0, tf.shape(x)[1], 1.0)[:, tf.newaxis] 108 | positional_embedding = tf.concat([ 109 | tf.math.sin(pos * inverse_frequency), 110 | tf.math.cos(pos * inverse_frequency)], axis=1)[tf.newaxis] 111 | 112 | # add the positional encoding 113 | h = tfkl.Add()([x, positional_embedding]) 114 | h = tfkl.LayerNormalization()(h) 115 | 116 | # first convolution layer in a residual block 117 | h = tfkl.Conv1D(hidden_size, kernel_size, 118 | padding='same', activation=None)(h) 119 | h = tfkl.LayerNormalization()(h) 120 | h = tfkl.Activation(activation)(h) 121 | 122 | # second convolution layer in a residual block 123 | h = tfkl.Conv1D(hidden_size, kernel_size, 124 | padding='same', activation=None)(h) 125 | h = tfkl.LayerNormalization()(h) 126 | h = tfkl.Activation(activation)(h) 127 | 128 | # add a residual connection to the model 129 | x = tfkl.Add()([x, h]) 130 | 131 | # flatten the result and predict the params of a gaussian 132 | logits = tfkl.Dense(task.num_classes)(x) 133 | 134 | # save the encoder as a keras model 135 | self.decoder_cnn = keras.Model( 136 | inputs=input_layer, outputs=logits) 137 | 138 | def encode(self, x_batch, training=False): 139 | mean, standard_dev = self.encoder_cnn(x_batch, training=training) 140 | return tfpd.MultivariateNormalDiag(loc=mean, scale_diag=standard_dev) 141 | 142 | def decode(self, z, training=False): 143 | logits = self.decoder_cnn(z, training=training) 144 | return tfpd.Categorical(logits=logits) 145 | 146 | def generate(self, z, training=False): 147 | logits = self.decoder_cnn(z, training=training) 148 | return tf.argmax(logits, axis=2, output_type=tf.int32) 149 | 150 | 151 | class ForwardModel(tf.keras.Sequential): 152 | """A Fully Connected Network with 2 trainable layers""" 153 | 154 | distribution = tfpd.Normal 155 | 156 | def __init__(self, input_shape, hidden_size=50, 157 | num_layers=1, initial_max_std=1.5, initial_min_std=0.5): 158 | """Create a fully connected architecture using keras that can process 159 | designs and predict a gaussian distribution over scores 160 | 161 | Args: 162 | 163 | task: StaticGraphTask 164 | a model-based optimization task 165 | embedding_size: int 166 | the size of the embedding matrix for discrete tasks 167 | hidden_size: int 168 | the global hidden size of the neural network 169 | num_layers: int 170 | the number of hidden layers 171 | initial_max_std: float 172 | the starting upper bound of the standard deviation 173 | initial_min_std: float 174 | the starting lower bound of the standard deviation 175 | 176 | """ 177 | 178 | self.max_logstd = tf.Variable(tf.fill([1, 1], np.log( 179 | initial_max_std).astype(np.float32)), trainable=True) 180 | self.min_logstd = tf.Variable(tf.fill([1, 1], np.log( 181 | initial_min_std).astype(np.float32)), trainable=True) 182 | 183 | layers = [] 184 | layers.append(tfkl.Flatten(input_shape=input_shape) 185 | if len(layers) == 0 else tfkl.Flatten()) 186 | for i in range(num_layers): 187 | layers.extend([tfkl.Dense(hidden_size), tfkl.LeakyReLU()]) 188 | 189 | layers.append(tfkl.Dense(2)) 190 | super(ForwardModel, self).__init__(layers) 191 | 192 | def get_params(self, inputs, **kwargs): 193 | """Return a dictionary of parameters for a particular distribution 194 | family such as the mean and variance of a gaussian 195 | 196 | Args: 197 | 198 | inputs: tf.Tensor 199 | a batch of training inputs shaped like [batch_size, channels] 200 | 201 | Returns: 202 | 203 | parameters: dict 204 | a dictionary that contains 'loc' and 'scale_diag' keys 205 | """ 206 | 207 | prediction = super(ForwardModel, self).__call__(inputs, **kwargs) 208 | mean, logstd = tf.split(prediction, 2, axis=-1) 209 | logstd = self.max_logstd - tf.nn.softplus(self.max_logstd - logstd) 210 | logstd = self.min_logstd + tf.nn.softplus(logstd - self.min_logstd) 211 | return {"loc": mean, "scale": tf.math.exp(logstd)} 212 | 213 | def get_distribution(self, inputs, **kwargs): 214 | """Return a distribution over the outputs of this model, for example 215 | a Multivariate Gaussian Distribution 216 | 217 | Args: 218 | 219 | inputs: tf.Tensor 220 | a batch of training inputs shaped like [batch_size, channels] 221 | 222 | Returns: 223 | 224 | distribution: tfp.distribution.Distribution 225 | a tensorflow probability distribution over outputs of the model 226 | """ 227 | 228 | return self.distribution(**self.get_params(inputs, **kwargs)) 229 | -------------------------------------------------------------------------------- /design_baselines/coms_cleaned/nets.py: -------------------------------------------------------------------------------- 1 | from tensorflow_probability import distributions as tfpd 2 | import tensorflow.keras.layers as tfkl 3 | import tensorflow.keras as keras 4 | import tensorflow as tf 5 | import numpy as np 6 | 7 | 8 | class TanhMultiplier(tf.keras.layers.Layer): 9 | 10 | def __init__(self, **kwargs): 11 | super(TanhMultiplier, self).__init__(**kwargs) 12 | w_init = tf.constant_initializer(1.0) 13 | self.multiplier = tf.Variable(initial_value=w_init( 14 | shape=(1,), dtype="float32"), trainable=True) 15 | 16 | def call(self, inputs, **kwargs): 17 | exp_multiplier = tf.math.exp(self.multiplier) 18 | return tf.math.tanh(inputs / exp_multiplier) * exp_multiplier 19 | 20 | 21 | def ForwardModel(input_shape, 22 | activations=('relu', 'relu'), 23 | hidden_size=2048, 24 | final_tanh=False): 25 | """Creates a tensorflow model that outputs a probability distribution 26 | specifying the score corresponding to an input x. 27 | 28 | Args: 29 | 30 | input_shape: tuple[int] 31 | the shape of input tensors to the model 32 | activations: tuple[str] 33 | the name of activation functions for every hidden layer 34 | hidden: int 35 | the global hidden size of the network 36 | max_std: float 37 | the upper bound of the learned standard deviation 38 | min_std: float 39 | the lower bound of the learned standard deviation 40 | """ 41 | 42 | activations = [tfkl.LeakyReLU if act == 'leaky_relu' else 43 | act for act in activations] 44 | 45 | layers = [tfkl.Flatten(input_shape=input_shape)] 46 | for act in activations: 47 | layers.extend([tfkl.Dense(hidden_size), tfkl.Activation(act) 48 | if isinstance(act, str) else act()]) 49 | layers.extend([tfkl.Dense(1)]) 50 | if final_tanh: 51 | layers.extend([TanhMultiplier()]) 52 | return tf.keras.Sequential(layers) 53 | 54 | 55 | class SequentialVAE(tf.keras.Model): 56 | 57 | def __init__(self, task, hidden_size=64, latent_size=256, 58 | activation='relu', kernel_size=3, num_blocks=4): 59 | 60 | super(SequentialVAE, self).__init__() 61 | 62 | input_shape = task.input_shape 63 | shape_before_flat = [input_shape[0] // ( 64 | 2 ** (num_blocks - 1)), hidden_size] 65 | 66 | """ 67 | 68 | DEFINE AN ENCODER MODEL THAT DOWNSAMPLES 69 | 70 | """ 71 | 72 | # the input layer of a keras model 73 | x = input_layer = keras.Input(shape=input_shape) 74 | 75 | # build a model with an input layer and optional embedding 76 | x = tfkl.Embedding(task.num_classes, hidden_size)(x) 77 | 78 | # the exponent of a positional embedding 79 | inverse_frequency = 1.0 / (10000.0 ** (tf.range( 80 | 0.0, hidden_size, 2.0) / hidden_size))[tf.newaxis] 81 | 82 | # calculate a positional embedding to break symmetry 83 | pos = tf.range(0.0, tf.shape(x)[1], 1.0)[:, tf.newaxis] 84 | positional_embedding = tf.concat([ 85 | tf.math.sin(pos * inverse_frequency), 86 | tf.math.cos(pos * inverse_frequency)], axis=1)[tf.newaxis] 87 | 88 | # add the positional encoding 89 | x = tfkl.Add()([x, positional_embedding]) 90 | x = tfkl.LayerNormalization()(x) 91 | 92 | # add several residual blocks to the model 93 | for i in range(num_blocks): 94 | 95 | if i > 0: 96 | # downsample the input sequence by 2 97 | x = tf.keras.layers.AveragePooling1D(pool_size=2, 98 | padding='same')(x) 99 | 100 | # first convolution layer in a residual block 101 | h = tfkl.Conv1D(hidden_size, kernel_size, 102 | padding='same', activation=None)(x) 103 | h = tfkl.LayerNormalization()(h) 104 | h = tfkl.Activation(activation)(h) 105 | 106 | # second convolution layer in a residual block 107 | h = tfkl.Conv1D(hidden_size, kernel_size, 108 | padding='same', activation=None)(h) 109 | h = tfkl.LayerNormalization()(h) 110 | h = tfkl.Activation(activation)(h) 111 | 112 | # add a residual connection to the model 113 | x = tfkl.Add()([x, h]) 114 | 115 | # flatten the result and predict the params of a gaussian 116 | flattened_x = tfkl.Flatten()(x) 117 | latent_mean = tfkl.Dense(latent_size)(flattened_x) 118 | latent_standard_dev = tfkl.Dense( 119 | latent_size, activation=tf.exp)(flattened_x) 120 | 121 | # save the encoder as a keras model 122 | self.encoder_cnn = keras.Model( 123 | inputs=input_layer, 124 | outputs=[latent_mean, latent_standard_dev]) 125 | 126 | """ 127 | 128 | DEFINE A DECODER THAT UPSAMPLES 129 | 130 | """ 131 | 132 | # the input layer of a keras model 133 | x = input_layer = keras.Input(shape=[latent_size]) 134 | x = tfkl.Dense(np.prod(shape_before_flat))(x) 135 | x = tfkl.Reshape(shape_before_flat)(x) 136 | 137 | # add several residual blocks to the model 138 | for i in reversed(range(num_blocks)): 139 | 140 | if i > 0: 141 | # up-sample the sequence 142 | x = tf.pad(tf.repeat(x, 2, axis=1), [[0, 0], [ 143 | 0, 144 | (input_shape[0] // (2 ** (i - 1))) % 2 145 | ], [0, 0]], mode="SYMMETRIC") 146 | 147 | # the exponent of a positional embedding 148 | inverse_frequency = 1.0 / (10000.0 ** (tf.range( 149 | 0.0, hidden_size, 2.0) / hidden_size))[tf.newaxis] 150 | 151 | # calculate a positional embedding to break symmetry 152 | pos = tf.range(0.0, tf.shape(x)[1], 1.0)[:, tf.newaxis] 153 | positional_embedding = tf.concat([ 154 | tf.math.sin(pos * inverse_frequency), 155 | tf.math.cos(pos * inverse_frequency)], axis=1)[tf.newaxis] 156 | 157 | # add the positional encoding 158 | h = tfkl.Add()([x, positional_embedding]) 159 | h = tfkl.LayerNormalization()(h) 160 | 161 | # first convolution layer in a residual block 162 | h = tfkl.Conv1D(hidden_size, kernel_size, 163 | padding='same', activation=None)(h) 164 | h = tfkl.LayerNormalization()(h) 165 | h = tfkl.Activation(activation)(h) 166 | 167 | # second convolution layer in a residual block 168 | h = tfkl.Conv1D(hidden_size, kernel_size, 169 | padding='same', activation=None)(h) 170 | h = tfkl.LayerNormalization()(h) 171 | h = tfkl.Activation(activation)(h) 172 | 173 | # add a residual connection to the model 174 | x = tfkl.Add()([x, h]) 175 | 176 | # flatten the result and predict the params of a gaussian 177 | logits = tfkl.Dense(task.num_classes)(x) 178 | 179 | # save the encoder as a keras model 180 | self.decoder_cnn = keras.Model( 181 | inputs=input_layer, outputs=logits) 182 | 183 | def encode(self, x_batch, training=False): 184 | mean, standard_dev = self.encoder_cnn(x_batch, training=training) 185 | return tfpd.MultivariateNormalDiag(loc=mean, scale_diag=standard_dev) 186 | 187 | def decode(self, z, training=False): 188 | logits = self.decoder_cnn(z, training=training) 189 | return tfpd.Categorical(logits=logits) 190 | 191 | def generate(self, z, training=False): 192 | logits = self.decoder_cnn(z, training=training) 193 | return tf.argmax(logits, axis=2, output_type=tf.int32) 194 | -------------------------------------------------------------------------------- /design_baselines/coms_original/__init__.py: -------------------------------------------------------------------------------- 1 | from design_baselines.data import StaticGraphTask 2 | from design_baselines.logger import Logger 3 | from design_baselines.utils import spearman 4 | from design_baselines.coms_original.trainers import ConservativeMaximumLikelihood 5 | from design_baselines.coms_original.trainers import TransformedMaximumLikelihood 6 | from design_baselines.coms_original.nets import ForwardModel 7 | from collections import defaultdict 8 | import tensorflow as tf 9 | import numpy as np 10 | import os 11 | 12 | 13 | def coms_original(config): 14 | """Train a forward model and perform offline model-based 15 | optimization using a conservative objective model 16 | 17 | Args: 18 | 19 | config: dict 20 | a dictionary of hyper parameters such as the learning rate 21 | """ 22 | 23 | # create the training task and logger 24 | logger = Logger(config['logging_dir']) 25 | task = StaticGraphTask(config['task'], **config['task_kwargs']) 26 | 27 | # save the initial dataset statistics for safe keeping 28 | x = task.x 29 | y = task.y 30 | 31 | # if the task is discrete then use a continuous relaxation 32 | if config['is_discrete']: 33 | p = np.full_like(x, 1 / float(x.shape[-1])) 34 | x = config.get('discrete_clip', 5.0) * x + ( 35 | 1.0 - config.get('discrete_clip', 5.0)) * p 36 | x = np.log(x) 37 | x = x[:, :, 1:] - x[:, :, :1] 38 | 39 | if config['normalize_ys']: 40 | # compute normalization statistics for the score 41 | mu_y = np.mean(y, axis=0, keepdims=True).astype(np.float32) 42 | y = y - mu_y 43 | st_y = np.std(y, axis=0, keepdims=True).astype(np.float32).clip(1e-6, 1e9) 44 | y = y / st_y 45 | else: 46 | # compute normalization statistics for the score 47 | mu_y = np.zeros_like(y[:1]) 48 | st_y = np.ones_like(y[:1]) 49 | 50 | if config['normalize_xs']: 51 | # compute normalization statistics for the data vectors 52 | mu_x = np.mean(x, axis=0, keepdims=True).astype(np.float32) 53 | x = x - mu_x 54 | st_x = np.std(x, axis=0, keepdims=True).astype(np.float32).clip(1e-6, 1e9) 55 | x = x / st_x 56 | else: 57 | # compute normalization statistics for the score 58 | mu_x = np.zeros_like(x[:1]) 59 | st_x = np.ones_like(x[:1]) 60 | 61 | input_shape = list(task.input_shape) 62 | if config['is_discrete']: 63 | input_shape[-1] = input_shape[-1] - 1 64 | 65 | solver_lr = config['solver_lr'] * np.sqrt(np.prod(input_shape)) 66 | solver_interval = int(config['solver_interval'] * ( 67 | x.shape[0] - config['val_size']) / config['batch_size']) 68 | solver_warmup = int(config['solver_warmup'] * ( 69 | x.shape[0] - config['val_size']) / config['batch_size']) 70 | 71 | # make a neural network to predict scores 72 | forward_model = ForwardModel( 73 | input_shape, 74 | activations=config['activations'], 75 | hidden=config['hidden_size'], 76 | final_tanh=config['final_tanh']) 77 | 78 | # create a trainer for a forward model with a conservative objective 79 | trainer = ConservativeMaximumLikelihood( 80 | forward_model, forward_model_opt=tf.keras.optimizers.Adam, 81 | forward_model_lr=config['forward_model_lr'], 82 | initial_alpha=config['initial_alpha'], 83 | alpha_opt=tf.keras.optimizers.Adam, 84 | alpha_lr=config['alpha_lr'], 85 | target_conservatism=config['target_conservatism'], 86 | negatives_fraction=config['negatives_fraction'], 87 | lookahead_steps=config['lookahead_steps'], 88 | lookahead_backprop=config['lookahead_backprop'], 89 | solver_beta=config['solver_beta'], 90 | solver_lr=solver_lr, 91 | solver_interval=solver_interval, 92 | solver_warmup=solver_warmup, 93 | solver_steps=config['solver_steps'], 94 | constraint_type=config['constraint_type'], 95 | entropy_coefficient=config['entropy_coefficient'], 96 | continuous_noise_std=config.get('continuous_noise_std', 0.0)) 97 | 98 | # make a neural network to predict scores 99 | validation_models = [ForwardModel( 100 | input_shape, 101 | activations=config['activations'], 102 | hidden=config['hidden_size'], 103 | final_tanh=config['final_tanh'])] 104 | 105 | # create a trainer for a forward model with a conservative objective 106 | validation_trainers = [TransformedMaximumLikelihood( 107 | model, 108 | forward_model_optim=tf.keras.optimizers.Adam, 109 | forward_model_lr=config['forward_model_lr'], 110 | continuous_noise_std=config.get('continuous_noise_std', 0.0), 111 | logger_prefix=f"validation_model_{i}") 112 | for i, model in enumerate(validation_models)] 113 | 114 | # create a data set 115 | train_data, validate_data = task.build( 116 | x=x, y=y, 117 | batch_size=config['batch_size'], 118 | val_size=config['val_size']) 119 | 120 | # train the validation models 121 | for t in validation_trainers: 122 | t.launch(train_data, validate_data, logger, 100) 123 | 124 | # select the top k initial designs from the dataset 125 | indices = tf.math.top_k(y[:, 0], k=config['batch_size'])[1] 126 | initial_x = tf.gather(x, indices, axis=0) 127 | 128 | # create the starting point for the optimizer 129 | evaluations = 0 130 | score = None 131 | trainer.solution = tf.Variable(initial_x) 132 | trainer.done = tf.Variable(tf.fill( 133 | [config['batch_size']] + [1 for _ in x.shape[1:]], False)) 134 | 135 | def evaluate_solution(xt): 136 | nonlocal evaluations, score 137 | 138 | # evaluate the design using the oracle and the forward model 139 | with tf.GradientTape() as tape: 140 | tape.watch(xt) 141 | model = forward_model(xt) 142 | 143 | # evaluate the predictions and gradient norm 144 | evaluations += 1 145 | grads = tape.gradient(model, xt) 146 | model = model * st_y + mu_y 147 | 148 | for i, val in enumerate(validation_models): 149 | prediction = val(xt) 150 | logger.record(f"validation_model_{i}/prediction", 151 | prediction * st_y + mu_y, evaluations) 152 | 153 | # record the prediction and score to the logger 154 | logger.record("distance/travelled", 155 | tf.linalg.norm(xt - initial_x), evaluations) 156 | logger.record(f"train/prediction", model, evaluations) 157 | logger.record(f"train/grad_norm", tf.linalg.norm( 158 | tf.reshape(grads, [grads.shape[0], -1]), axis=-1), evaluations) 159 | 160 | if evaluations in config['evaluate_steps'] \ 161 | or len(config['evaluate_steps']) == 0 or score is None: 162 | solution = xt * st_x + mu_x 163 | if config['is_discrete']: 164 | solution = tf.math.softmax( 165 | tf.pad(solution, [[0, 0], [0, 0], [1, 0]]) / 0.001) 166 | score = task.score(solution) 167 | logger.record("score", score, evaluations, percentile=True) 168 | logger.record(f"rank_corr/model_to_real", 169 | spearman(model[:, 0], score[:, 0]), evaluations) 170 | 171 | return score, model 172 | 173 | # keep track of when to record performance 174 | interval = trainer.solver_interval 175 | warmup = trainer.solver_warmup 176 | 177 | scores = [] 178 | predictions = [] 179 | 180 | # train model for many epochs with conservatism 181 | for e in range(config['epochs']): 182 | 183 | statistics = defaultdict(list) 184 | for x, y in train_data: 185 | for name, tensor in trainer.train_step(x, y).items(): 186 | statistics[name].append(tensor) 187 | 188 | # evaluate the current solution 189 | if tf.logical_and( 190 | tf.equal(tf.math.mod(trainer.step, interval), 0), 191 | tf.math.greater_equal(trainer.step, warmup)): 192 | score, model = evaluate_solution(trainer.solution) 193 | scores.append(score) 194 | predictions.append(model.numpy()) 195 | 196 | for name in statistics.keys(): 197 | logger.record( 198 | name, tf.concat(statistics[name], axis=0), e) 199 | 200 | statistics = defaultdict(list) 201 | for x, y in validate_data: 202 | for name, tensor in trainer.validate_step(x, y).items(): 203 | statistics[name].append(tensor) 204 | 205 | for name in statistics.keys(): 206 | logger.record( 207 | name, tf.concat(statistics[name], axis=0), e) 208 | 209 | if tf.reduce_all(trainer.done): 210 | break 211 | 212 | # save the model predictions and scores to be aggregated later 213 | np.save(os.path.join(config['logging_dir'], "scores.npy"), 214 | np.concatenate(scores, axis=1)) 215 | np.save(os.path.join(config['logging_dir'], "predictions.npy"), 216 | np.stack(predictions, axis=1)) 217 | -------------------------------------------------------------------------------- /design_baselines/coms_original/nets.py: -------------------------------------------------------------------------------- 1 | import tensorflow.keras.layers as tfkl 2 | import tensorflow as tf 3 | 4 | 5 | class TanhMultiplier(tf.keras.layers.Layer): 6 | 7 | def __init__(self, **kwargs): 8 | super(TanhMultiplier, self).__init__(**kwargs) 9 | w_init = tf.constant_initializer(1.0) 10 | self.multiplier = tf.Variable(initial_value=w_init( 11 | shape=(1,), dtype="float32"), trainable=True) 12 | 13 | def call(self, inputs, **kwargs): 14 | exp_multiplier = tf.math.exp(self.multiplier) 15 | return tf.math.tanh(inputs / exp_multiplier) * exp_multiplier 16 | 17 | 18 | def ForwardModel(input_shape, 19 | activations=('relu', 'relu'), 20 | hidden=2048, 21 | final_tanh=False): 22 | """Creates a tensorflow model that outputs a probability distribution 23 | specifying the score corresponding to an input x. 24 | 25 | Args: 26 | 27 | input_shape: tuple[int] 28 | the shape of input tensors to the model 29 | activations: tuple[str] 30 | the name of activation functions for every hidden layer 31 | hidden: int 32 | the global hidden size of the network 33 | max_std: float 34 | the upper bound of the learned standard deviation 35 | min_std: float 36 | the lower bound of the learned standard deviation 37 | """ 38 | 39 | activations = [tfkl.LeakyReLU if act == 'leaky_relu' else 40 | tfkl.Activation(tf.math.cos) if act == 'cos' else 41 | act for act in activations] 42 | 43 | layers = [tfkl.Flatten(input_shape=input_shape)] 44 | for act in activations: 45 | layers.extend([tfkl.Dense(hidden), tfkl.Activation(act) 46 | if isinstance(act, str) else act()]) 47 | layers.extend([tfkl.Dense(1)]) 48 | if final_tanh: 49 | layers.extend([TanhMultiplier()]) 50 | return tf.keras.Sequential(layers) 51 | -------------------------------------------------------------------------------- /design_baselines/gradient_ascent/__init__.py: -------------------------------------------------------------------------------- 1 | from design_baselines.data import StaticGraphTask, build_pipeline 2 | from design_baselines.logger import Logger 3 | from design_baselines.utils import spearman 4 | from design_baselines.gradient_ascent.trainers import MaximumLikelihood 5 | from design_baselines.gradient_ascent.trainers import Ensemble, VAETrainer 6 | from design_baselines.gradient_ascent.nets import ForwardModel, SequentialVAE 7 | import tensorflow_probability as tfp 8 | import tensorflow as tf 9 | import numpy as np 10 | import os 11 | 12 | 13 | def gradient_ascent(config): 14 | """Train a Score Function to solve a Model-Based Optimization 15 | using gradient ascent on the input design 16 | 17 | Args: 18 | 19 | config: dict 20 | a dictionary of hyper parameters such as the learning rate 21 | """ 22 | 23 | # create the training task and logger 24 | logger = Logger(config['logging_dir']) 25 | task = StaticGraphTask(config['task'], **config['task_kwargs']) 26 | 27 | if config['normalize_ys']: 28 | task.map_normalize_y() 29 | if task.is_discrete and not config["use_vae"]: 30 | task.map_to_logits() 31 | if config['normalize_xs']: 32 | task.map_normalize_x() 33 | 34 | x = task.x 35 | y = task.y 36 | 37 | if task.is_discrete and config["use_vae"]: 38 | 39 | vae_model = SequentialVAE( 40 | task, 41 | hidden_size=config['vae_hidden_size'], 42 | latent_size=config['vae_latent_size'], 43 | activation=config['vae_activation'], 44 | kernel_size=config['vae_kernel_size'], 45 | num_blocks=config['vae_num_blocks']) 46 | 47 | vae_trainer = VAETrainer(vae_model, 48 | vae_optim=tf.keras.optimizers.Adam, 49 | vae_lr=config['vae_lr'], 50 | beta=config['vae_beta']) 51 | 52 | # create the training task and logger 53 | train_data, val_data = build_pipeline( 54 | x=x, y=y, 55 | batch_size=config['vae_batch_size'], 56 | val_size=config['val_size']) 57 | 58 | # estimate the number of training steps per epoch 59 | vae_trainer.launch(train_data, val_data, 60 | logger, config['vae_epochs']) 61 | 62 | # map the x values to latent space 63 | x = vae_model.encoder_cnn.predict(x)[0] 64 | 65 | mean = np.mean(x, axis=0, keepdims=True) 66 | standard_dev = np.std(x - mean, axis=0, keepdims=True) 67 | x = (x - mean) / standard_dev 68 | 69 | input_shape = x.shape[1:] 70 | input_size = np.prod(input_shape) 71 | 72 | # make several keras neural networks with different architectures 73 | forward_models = [ForwardModel( 74 | input_shape, 75 | activations=activations, 76 | hidden_size=config['hidden_size'], 77 | initial_max_std=config['initial_max_std'], 78 | initial_min_std=config['initial_min_std']) 79 | for activations in config['activations']] 80 | 81 | # scale the learning rate based on the number of channels in x 82 | config['solver_lr'] *= np.sqrt(np.prod(x.shape[1:])) 83 | 84 | trs = [] 85 | for i, fm in enumerate(forward_models): 86 | 87 | # create a bootstrapped data set 88 | train_data, validate_data = build_pipeline( 89 | x=x, y=y, batch_size=config['batch_size'], 90 | val_size=config['val_size'], bootstraps=1) 91 | 92 | # create a trainer for a forward model with a conservative objective 93 | trainer = MaximumLikelihood( 94 | fm, 95 | forward_model_optim=tf.keras.optimizers.Adam, 96 | forward_model_lr=config['forward_model_lr'], 97 | noise_std=config.get('model_noise_std', 0.0)) 98 | 99 | # train the model for an additional number of epochs 100 | trs.append(trainer) 101 | trainer.launch(train_data, validate_data, logger, 102 | config['epochs'], header=f'oracle_{i}/') 103 | 104 | # select the top k initial designs from the dataset 105 | mean_x = tf.reduce_mean(x, axis=0, keepdims=True) 106 | indices = tf.math.top_k(y[:, 0], k=config['solver_samples'])[1] 107 | initial_x = tf.gather(x, indices, axis=0) 108 | x = initial_x 109 | 110 | # evaluate the starting point 111 | solution = x 112 | if task.is_normalized_y: 113 | preds = [task.denormalize_y( 114 | fm.get_distribution(solution).mean()) 115 | for fm in forward_models] 116 | else: 117 | preds = [fm.get_distribution(solution).mean() 118 | for fm in forward_models] 119 | 120 | # record the prediction and score to the logger 121 | logger.record("distance/travelled", tf.linalg.norm(solution - initial_x), 0) 122 | logger.record("distance/from_mean", tf.linalg.norm(solution - mean_x), 0) 123 | for n, prediction_i in enumerate(preds): 124 | logger.record(f"oracle_{n}/prediction", prediction_i, 0) 125 | if n > 0: 126 | logger.record(f"rank_corr/0_to_{n}", 127 | spearman(preds[0][:, 0], prediction_i[:, 0]), 0) 128 | 129 | # perform gradient ascent on the score through the forward model 130 | for i in range(1, config['solver_steps'] + 1): 131 | # back propagate through the forward model 132 | with tf.GradientTape() as tape: 133 | tape.watch(x) 134 | predictions = [] 135 | for fm in forward_models: 136 | solution = x 137 | predictions.append(fm.get_distribution(solution).mean()) 138 | if config['aggregation_method'] == 'mean': 139 | score = tf.reduce_min(predictions, axis=0) 140 | if config['aggregation_method'] == 'min': 141 | score = tf.reduce_min(predictions, axis=0) 142 | if config['aggregation_method'] == 'random': 143 | score = predictions[np.random.randint(len(predictions))] 144 | grads = tape.gradient(score, x) 145 | 146 | # use the conservative optimizer to update the solution 147 | x = x + config['solver_lr'] * grads 148 | solution = x 149 | 150 | # evaluate the design using the oracle and the forward model 151 | if task.is_normalized_y: 152 | preds = [task.denormalize_y( 153 | fm.get_distribution(solution).mean()) 154 | for fm in forward_models] 155 | else: 156 | preds = [fm.get_distribution(solution).mean() 157 | for fm in forward_models] 158 | 159 | # record the prediction and score to the logger 160 | logger.record("distance/travelled", tf.linalg.norm(solution - initial_x), i) 161 | logger.record("distance/from_mean", tf.linalg.norm(solution - mean_x), i) 162 | for n, prediction_i in enumerate(preds): 163 | logger.record(f"oracle_{n}/prediction", prediction_i, i) 164 | logger.record(f"oracle_{n}/grad_norm", tf.linalg.norm( 165 | tf.reshape(grads[n], [-1, input_size]), axis=-1), i) 166 | if n > 0: 167 | logger.record(f"rank_corr/0_to_{n}", 168 | spearman(preds[0][:, 0], prediction_i[:, 0]), i) 169 | logger.record(f"grad_corr/0_to_{n}", tfp.stats.correlation( 170 | grads[0], grads[n], sample_axis=0, event_axis=None), i) 171 | 172 | if task.is_discrete and config["use_vae"]: 173 | solution = solution * standard_dev + mean 174 | logits = vae_model.decoder_cnn.predict(solution) 175 | solution = tf.argmax(logits, axis=2, output_type=tf.int32) 176 | 177 | # save the current solution to the disk 178 | np.save(os.path.join(config["logging_dir"], 179 | f"solution.npy"), solution.numpy()) 180 | if config["do_evaluation"]: 181 | 182 | # evaluate the found solution and record a video 183 | score = task.predict(solution) 184 | if task.is_normalized_y: 185 | score = task.denormalize_y(score) 186 | logger.record("score", score, config['solver_steps'], percentile=True) 187 | -------------------------------------------------------------------------------- /design_baselines/gradient_ascent/distribution_experiments.py: -------------------------------------------------------------------------------- 1 | from ray import tune 2 | import click 3 | import ray 4 | import os 5 | 6 | 7 | @click.group() 8 | def cli(): 9 | """A group of experiments for training Conservative Score Models 10 | and reproducing our ICLR 2021 results. 11 | """ 12 | 13 | 14 | ############# 15 | 16 | 17 | @cli.command() 18 | @click.option('--local-dir', type=str, default='gradient-ascent-hopper') 19 | @click.option('--cpus', type=int, default=24) 20 | @click.option('--gpus', type=int, default=1) 21 | @click.option('--num-parallel', type=int, default=1) 22 | @click.option('--num-samples', type=int, default=1) 23 | def hopper(local_dir, cpus, gpus, num_parallel, num_samples): 24 | """Evaluate Naive Gradient Ascent on HopperController-Exact-v0 25 | """ 26 | 27 | # Final Version 28 | 29 | from design_baselines.gradient_ascent import gradient_ascent 30 | ray.init(num_cpus=cpus, 31 | num_gpus=gpus, 32 | include_dashboard=False, 33 | _temp_dir=os.path.expanduser('~/tmp')) 34 | tune.run(gradient_ascent, config={ 35 | "logging_dir": "data", 36 | "task": "HopperController-Exact-v0", 37 | "task_kwargs": {"relabel": False, "dataset_kwargs": { 38 | "max_samples": 1000, 39 | "distribution": tune.grid_search([ 40 | "uniform", 41 | "linear", 42 | "quadratic", 43 | "exponential", 44 | "circular" 45 | ]), 46 | "max_percentile": 100, 47 | "min_percentile": 0 48 | }}, 49 | "normalize_ys": True, 50 | "normalize_xs": True, 51 | "model_noise_std": 0.0, 52 | "val_size": 200, 53 | "batch_size": 128, 54 | "epochs": 100, 55 | "activations": [['leaky_relu', 'leaky_relu']], 56 | "hidden_size": 2048, 57 | "initial_max_std": 0.2, 58 | "initial_min_std": 0.1, 59 | "forward_model_lr": 0.0003, 60 | "aggregation_method": 'mean', 61 | "solver_samples": 512, 62 | "solver_lr": 0.01, 63 | "solver_steps": 200}, 64 | num_samples=num_samples, 65 | local_dir=local_dir, 66 | resources_per_trial={'cpu': cpus // num_parallel, 67 | 'gpu': gpus / num_parallel - 0.01}) 68 | 69 | 70 | @cli.command() 71 | @click.option('--local-dir', type=str, default='gradient-ascent-superconductor') 72 | @click.option('--cpus', type=int, default=24) 73 | @click.option('--gpus', type=int, default=1) 74 | @click.option('--num-parallel', type=int, default=1) 75 | @click.option('--num-samples', type=int, default=1) 76 | @click.option('--oracle', type=str, default="RandomForest") 77 | def superconductor(local_dir, cpus, gpus, num_parallel, num_samples, oracle): 78 | """Evaluate Naive Gradient Ascent on Superconductor-RandomForest-v0 79 | """ 80 | 81 | # Final Version 82 | 83 | from design_baselines.gradient_ascent import gradient_ascent 84 | ray.init(num_cpus=cpus, 85 | num_gpus=gpus, 86 | include_dashboard=False, 87 | _temp_dir=os.path.expanduser('~/tmp')) 88 | tune.run(gradient_ascent, config={ 89 | "logging_dir": "data", 90 | "task": f"Superconductor-{oracle}-v0", 91 | "task_kwargs": {"relabel": False, "dataset_kwargs": { 92 | "max_samples": 5000, 93 | "distribution": tune.grid_search([ 94 | "uniform", 95 | "linear", 96 | "quadratic", 97 | "exponential", 98 | "circular" 99 | ]), 100 | "max_percentile": 100, 101 | "min_percentile": 0 102 | }}, 103 | "normalize_ys": True, 104 | "normalize_xs": True, 105 | "model_noise_std": 0.0, 106 | "val_size": 200, 107 | "batch_size": 128, 108 | "epochs": 50, 109 | "activations": [['leaky_relu', 'leaky_relu']], 110 | "hidden_size": 2048, 111 | "initial_max_std": 0.2, 112 | "initial_min_std": 0.1, 113 | "forward_model_lr": 0.0003, 114 | "aggregation_method": 'mean', 115 | "solver_samples": 512, 116 | "solver_lr": 0.01, 117 | "solver_steps": 200}, 118 | num_samples=num_samples, 119 | local_dir=local_dir, 120 | resources_per_trial={'cpu': cpus // num_parallel, 121 | 'gpu': gpus / num_parallel - 0.01}) 122 | 123 | 124 | @cli.command() 125 | @click.option('--local-dir', type=str, default='gradient-ascent-gfp') 126 | @click.option('--cpus', type=int, default=24) 127 | @click.option('--gpus', type=int, default=1) 128 | @click.option('--num-parallel', type=int, default=1) 129 | @click.option('--num-samples', type=int, default=1) 130 | @click.option('--oracle', type=str, default="Transformer") 131 | def gfp(local_dir, cpus, gpus, num_parallel, num_samples, oracle): 132 | """Evaluate Naive Gradient Ascent on GFP-Transformer-v0 133 | """ 134 | 135 | # Final Version 136 | 137 | from design_baselines.gradient_ascent import gradient_ascent 138 | ray.init(num_cpus=cpus, 139 | num_gpus=gpus, 140 | include_dashboard=False, 141 | _temp_dir=os.path.expanduser('~/tmp')) 142 | tune.run(gradient_ascent, config={ 143 | "logging_dir": "data", 144 | "task": f"GFP-{oracle}-v0", 145 | "task_kwargs": {"relabel": False, "dataset_kwargs": { 146 | "max_samples": 5000, 147 | "distribution": tune.grid_search([ 148 | "uniform", 149 | "linear", 150 | "quadratic", 151 | "exponential", 152 | "circular" 153 | ]), 154 | "max_percentile": 100, 155 | "min_percentile": 0 156 | }}, 157 | "normalize_ys": True, 158 | "normalize_xs": True, 159 | "model_noise_std": 0.0, 160 | "val_size": 200, 161 | "use_vae": False, 162 | "vae_beta": 0.005, 163 | "vae_epochs": 50, 164 | "vae_batch_size": 128, 165 | "vae_hidden_size": 64, 166 | "vae_latent_size": 256, 167 | "vae_activation": "relu", 168 | "vae_kernel_size": 3, 169 | "vae_num_blocks": 5, 170 | "vae_lr": 0.0003, 171 | "batch_size": 128, 172 | "epochs": 100, 173 | "activations": [['leaky_relu', 'leaky_relu']], 174 | "hidden_size": 2048, 175 | "initial_max_std": 0.2, 176 | "initial_min_std": 0.1, 177 | "forward_model_lr": 0.0003, 178 | "aggregation_method": 'mean', 179 | "solver_samples": 512, 180 | "solver_lr": 0.01, 181 | "solver_steps": 200}, 182 | num_samples=num_samples, 183 | local_dir=local_dir, 184 | resources_per_trial={'cpu': cpus // num_parallel, 185 | 'gpu': gpus / num_parallel - 0.01}) 186 | 187 | 188 | @cli.command() 189 | @click.option('--local-dir', type=str, default='gradient-ascent-utr') 190 | @click.option('--cpus', type=int, default=24) 191 | @click.option('--gpus', type=int, default=1) 192 | @click.option('--num-parallel', type=int, default=1) 193 | @click.option('--num-samples', type=int, default=1) 194 | def utr(local_dir, cpus, gpus, num_parallel, num_samples): 195 | """Evaluate Naive Gradient Ascent on UTR-ResNet-v0 196 | """ 197 | 198 | # Final Version 199 | 200 | from design_baselines.gradient_ascent import gradient_ascent 201 | ray.init(num_cpus=cpus, 202 | num_gpus=gpus, 203 | include_dashboard=False, 204 | _temp_dir=os.path.expanduser('~/tmp')) 205 | tune.run(gradient_ascent, config={ 206 | "logging_dir": "data", 207 | "task": "UTR-ResNet-v0", 208 | "task_kwargs": {"relabel": False, "dataset_kwargs": { 209 | "max_samples": 20000, 210 | "distribution": tune.grid_search([ 211 | "uniform", 212 | "linear", 213 | "quadratic", 214 | "exponential", 215 | "circular" 216 | ]), 217 | "max_percentile": 100, 218 | "min_percentile": 0 219 | }}, 220 | "normalize_ys": True, 221 | "normalize_xs": True, 222 | "model_noise_std": 0.0, 223 | "val_size": 200, 224 | "use_vae": False, 225 | "vae_beta": 0.01, 226 | "vae_epochs": 50, 227 | "vae_batch_size": 128, 228 | "vae_hidden_size": 64, 229 | "vae_latent_size": 256, 230 | "vae_activation": "relu", 231 | "vae_kernel_size": 3, 232 | "vae_num_blocks": 5, 233 | "vae_lr": 0.0003, 234 | "batch_size": 128, 235 | "epochs": 100, 236 | "activations": [['leaky_relu', 'leaky_relu']], 237 | "hidden_size": 2048, 238 | "initial_max_std": 0.2, 239 | "initial_min_std": 0.1, 240 | "forward_model_lr": 0.0003, 241 | "aggregation_method": 'mean', 242 | "solver_samples": 512, 243 | "solver_lr": 0.01, 244 | "solver_steps": 200}, 245 | num_samples=num_samples, 246 | local_dir=local_dir, 247 | resources_per_trial={'cpu': cpus // num_parallel, 248 | 'gpu': gpus / num_parallel - 0.01}) 249 | -------------------------------------------------------------------------------- /design_baselines/gradient_ascent/distribution_mean_ensemble_experiments.py: -------------------------------------------------------------------------------- 1 | from ray import tune 2 | import click 3 | import ray 4 | import os 5 | 6 | 7 | @click.group() 8 | def cli(): 9 | """A group of experiments for training Conservative Score Models 10 | and reproducing our ICLR 2021 results. 11 | """ 12 | 13 | 14 | ############# 15 | 16 | 17 | @cli.command() 18 | @click.option('--local-dir', type=str, default='gradient-ascent-hopper') 19 | @click.option('--cpus', type=int, default=24) 20 | @click.option('--gpus', type=int, default=1) 21 | @click.option('--num-parallel', type=int, default=1) 22 | @click.option('--num-samples', type=int, default=1) 23 | def hopper(local_dir, cpus, gpus, num_parallel, num_samples): 24 | """Evaluate Naive Gradient Ascent on HopperController-Exact-v0 25 | """ 26 | 27 | # Final Version 28 | 29 | from design_baselines.gradient_ascent import gradient_ascent 30 | ray.init(num_cpus=cpus, 31 | num_gpus=gpus, 32 | include_dashboard=False, 33 | _temp_dir=os.path.expanduser('~/tmp')) 34 | tune.run(gradient_ascent, config={ 35 | "logging_dir": "data", 36 | "task": "HopperController-Exact-v0", 37 | "task_kwargs": {"relabel": False, "dataset_kwargs": { 38 | "max_samples": 1000, 39 | "distribution": tune.grid_search([ 40 | "uniform", 41 | "linear", 42 | "quadratic", 43 | "exponential", 44 | "circular" 45 | ]), 46 | "max_percentile": 100, 47 | "min_percentile": 0 48 | }}, 49 | "normalize_ys": True, 50 | "normalize_xs": True, 51 | "model_noise_std": 0.0, 52 | "val_size": 200, 53 | "batch_size": 128, 54 | "epochs": 100, 55 | "activations": [['leaky_relu', 'leaky_relu']] * 5, 56 | "hidden_size": 2048, 57 | "initial_max_std": 0.2, 58 | "initial_min_std": 0.1, 59 | "forward_model_lr": 0.0003, 60 | "aggregation_method": 'mean', 61 | "solver_samples": 512, 62 | "solver_lr": 0.01, 63 | "solver_steps": 200}, 64 | num_samples=num_samples, 65 | local_dir=local_dir, 66 | resources_per_trial={'cpu': cpus // num_parallel, 67 | 'gpu': gpus / num_parallel - 0.01}) 68 | 69 | 70 | @cli.command() 71 | @click.option('--local-dir', type=str, default='gradient-ascent-superconductor') 72 | @click.option('--cpus', type=int, default=24) 73 | @click.option('--gpus', type=int, default=1) 74 | @click.option('--num-parallel', type=int, default=1) 75 | @click.option('--num-samples', type=int, default=1) 76 | @click.option('--oracle', type=str, default="RandomForest") 77 | def superconductor(local_dir, cpus, gpus, num_parallel, num_samples, oracle): 78 | """Evaluate Naive Gradient Ascent on Superconductor-RandomForest-v0 79 | """ 80 | 81 | # Final Version 82 | 83 | from design_baselines.gradient_ascent import gradient_ascent 84 | ray.init(num_cpus=cpus, 85 | num_gpus=gpus, 86 | include_dashboard=False, 87 | _temp_dir=os.path.expanduser('~/tmp')) 88 | tune.run(gradient_ascent, config={ 89 | "logging_dir": "data", 90 | "task": f"Superconductor-{oracle}-v0", 91 | "task_kwargs": {"relabel": False, "dataset_kwargs": { 92 | "max_samples": 5000, 93 | "distribution": tune.grid_search([ 94 | "uniform", 95 | "linear", 96 | "quadratic", 97 | "exponential", 98 | "circular" 99 | ]), 100 | "max_percentile": 100, 101 | "min_percentile": 0 102 | }}, 103 | "normalize_ys": True, 104 | "normalize_xs": True, 105 | "model_noise_std": 0.0, 106 | "val_size": 200, 107 | "batch_size": 128, 108 | "epochs": 100, 109 | "activations": [['leaky_relu', 'leaky_relu']] * 5, 110 | "hidden_size": 2048, 111 | "initial_max_std": 0.2, 112 | "initial_min_std": 0.1, 113 | "forward_model_lr": 0.0003, 114 | "aggregation_method": 'mean', 115 | "solver_samples": 512, 116 | "solver_lr": 0.01, 117 | "solver_steps": 200}, 118 | num_samples=num_samples, 119 | local_dir=local_dir, 120 | resources_per_trial={'cpu': cpus // num_parallel, 121 | 'gpu': gpus / num_parallel - 0.01}) 122 | 123 | 124 | @cli.command() 125 | @click.option('--local-dir', type=str, default='gradient-ascent-gfp') 126 | @click.option('--cpus', type=int, default=24) 127 | @click.option('--gpus', type=int, default=1) 128 | @click.option('--num-parallel', type=int, default=1) 129 | @click.option('--num-samples', type=int, default=1) 130 | @click.option('--oracle', type=str, default="Transformer") 131 | def gfp(local_dir, cpus, gpus, num_parallel, num_samples, oracle): 132 | """Evaluate Naive Gradient Ascent on GFP-Transformer-v0 133 | """ 134 | 135 | # Final Version 136 | 137 | from design_baselines.gradient_ascent import gradient_ascent 138 | ray.init(num_cpus=cpus, 139 | num_gpus=gpus, 140 | include_dashboard=False, 141 | _temp_dir=os.path.expanduser('~/tmp')) 142 | tune.run(gradient_ascent, config={ 143 | "logging_dir": "data", 144 | "task": f"GFP-{oracle}-v0", 145 | "task_kwargs": {"relabel": False, "dataset_kwargs": { 146 | "max_samples": 5000, 147 | "distribution": tune.grid_search([ 148 | "uniform", 149 | "linear", 150 | "quadratic", 151 | "exponential", 152 | "circular" 153 | ]), 154 | "max_percentile": 100, 155 | "min_percentile": 0 156 | }}, 157 | "normalize_ys": True, 158 | "normalize_xs": True, 159 | "model_noise_std": 0.0, 160 | "val_size": 200, 161 | "use_vae": False, 162 | "vae_beta": 0.01, 163 | "vae_epochs": 50, 164 | "vae_batch_size": 128, 165 | "vae_hidden_size": 64, 166 | "vae_latent_size": 256, 167 | "vae_activation": "relu", 168 | "vae_kernel_size": 3, 169 | "vae_num_blocks": 5, 170 | "vae_lr": 0.0003, 171 | "batch_size": 128, 172 | "epochs": 100, 173 | "activations": [['leaky_relu', 'leaky_relu']] * 5, 174 | "hidden_size": 2048, 175 | "initial_max_std": 0.2, 176 | "initial_min_std": 0.1, 177 | "forward_model_lr": 0.0003, 178 | "aggregation_method": 'mean', 179 | "solver_samples": 512, 180 | "solver_lr": 0.01, 181 | "solver_steps": 200}, 182 | num_samples=num_samples, 183 | local_dir=local_dir, 184 | resources_per_trial={'cpu': cpus // num_parallel, 185 | 'gpu': gpus / num_parallel - 0.01}) 186 | 187 | 188 | @cli.command() 189 | @click.option('--local-dir', type=str, default='gradient-ascent-utr') 190 | @click.option('--cpus', type=int, default=24) 191 | @click.option('--gpus', type=int, default=1) 192 | @click.option('--num-parallel', type=int, default=1) 193 | @click.option('--num-samples', type=int, default=1) 194 | def utr(local_dir, cpus, gpus, num_parallel, num_samples): 195 | """Evaluate Naive Gradient Ascent on UTR-ResNet-v0 196 | """ 197 | 198 | # Final Version 199 | 200 | from design_baselines.gradient_ascent import gradient_ascent 201 | ray.init(num_cpus=cpus, 202 | num_gpus=gpus, 203 | include_dashboard=False, 204 | _temp_dir=os.path.expanduser('~/tmp')) 205 | tune.run(gradient_ascent, config={ 206 | "logging_dir": "data", 207 | "task": "UTR-ResNet-v0", 208 | "task_kwargs": {"relabel": False, "dataset_kwargs": { 209 | "max_samples": 20000, 210 | "distribution": tune.grid_search([ 211 | "uniform", 212 | "linear", 213 | "quadratic", 214 | "exponential", 215 | "circular" 216 | ]), 217 | "max_percentile": 100, 218 | "min_percentile": 0 219 | }}, 220 | "normalize_ys": True, 221 | "normalize_xs": True, 222 | "model_noise_std": 0.0, 223 | "val_size": 200, 224 | "use_vae": False, 225 | "vae_beta": 0.01, 226 | "vae_epochs": 50, 227 | "vae_batch_size": 128, 228 | "vae_hidden_size": 64, 229 | "vae_latent_size": 256, 230 | "vae_activation": "relu", 231 | "vae_kernel_size": 3, 232 | "vae_num_blocks": 5, 233 | "vae_lr": 0.0003, 234 | "batch_size": 128, 235 | "epochs": 100, 236 | "activations": [['leaky_relu', 'leaky_relu']] * 5, 237 | "hidden_size": 2048, 238 | "initial_max_std": 0.2, 239 | "initial_min_std": 0.1, 240 | "forward_model_lr": 0.0003, 241 | "aggregation_method": 'mean', 242 | "solver_samples": 512, 243 | "solver_lr": 0.01, 244 | "solver_steps": 200}, 245 | num_samples=num_samples, 246 | local_dir=local_dir, 247 | resources_per_trial={'cpu': cpus // num_parallel, 248 | 'gpu': gpus / num_parallel - 0.01}) 249 | -------------------------------------------------------------------------------- /design_baselines/gradient_ascent/distribution_min_ensemble_experiments.py: -------------------------------------------------------------------------------- 1 | from ray import tune 2 | import click 3 | import ray 4 | import os 5 | 6 | 7 | @click.group() 8 | def cli(): 9 | """A group of experiments for training Conservative Score Models 10 | and reproducing our ICLR 2021 results. 11 | """ 12 | 13 | 14 | ############# 15 | 16 | 17 | @cli.command() 18 | @click.option('--local-dir', type=str, default='gradient-ascent-hopper') 19 | @click.option('--cpus', type=int, default=24) 20 | @click.option('--gpus', type=int, default=1) 21 | @click.option('--num-parallel', type=int, default=1) 22 | @click.option('--num-samples', type=int, default=1) 23 | def hopper(local_dir, cpus, gpus, num_parallel, num_samples): 24 | """Evaluate Naive Gradient Ascent on HopperController-Exact-v0 25 | """ 26 | 27 | # Final Version 28 | 29 | from design_baselines.gradient_ascent import gradient_ascent 30 | ray.init(num_cpus=cpus, 31 | num_gpus=gpus, 32 | include_dashboard=False, 33 | _temp_dir=os.path.expanduser('~/tmp')) 34 | tune.run(gradient_ascent, config={ 35 | "logging_dir": "data", 36 | "task": "HopperController-Exact-v0", 37 | "task_kwargs": {"relabel": False, "dataset_kwargs": { 38 | "max_samples": 1000, 39 | "distribution": tune.grid_search([ 40 | "uniform", 41 | "linear", 42 | "quadratic", 43 | "exponential", 44 | "circular" 45 | ]), 46 | "max_percentile": 100, 47 | "min_percentile": 0 48 | }}, 49 | "normalize_ys": True, 50 | "normalize_xs": True, 51 | "model_noise_std": 0.0, 52 | "val_size": 200, 53 | "batch_size": 128, 54 | "epochs": 100, 55 | "activations": [['leaky_relu', 'leaky_relu']] * 5, 56 | "hidden_size": 2048, 57 | "initial_max_std": 0.2, 58 | "initial_min_std": 0.1, 59 | "forward_model_lr": 0.0003, 60 | "aggregation_method": 'min', 61 | "solver_samples": 512, 62 | "solver_lr": 0.01, 63 | "solver_steps": 200}, 64 | num_samples=num_samples, 65 | local_dir=local_dir, 66 | resources_per_trial={'cpu': cpus // num_parallel, 67 | 'gpu': gpus / num_parallel - 0.01}) 68 | 69 | 70 | @cli.command() 71 | @click.option('--local-dir', type=str, default='gradient-ascent-superconductor') 72 | @click.option('--cpus', type=int, default=24) 73 | @click.option('--gpus', type=int, default=1) 74 | @click.option('--num-parallel', type=int, default=1) 75 | @click.option('--num-samples', type=int, default=1) 76 | @click.option('--oracle', type=str, default="RandomForest") 77 | def superconductor(local_dir, cpus, gpus, num_parallel, num_samples, oracle): 78 | """Evaluate Naive Gradient Ascent on Superconductor-RandomForest-v0 79 | """ 80 | 81 | # Final Version 82 | 83 | from design_baselines.gradient_ascent import gradient_ascent 84 | ray.init(num_cpus=cpus, 85 | num_gpus=gpus, 86 | include_dashboard=False, 87 | _temp_dir=os.path.expanduser('~/tmp')) 88 | tune.run(gradient_ascent, config={ 89 | "logging_dir": "data", 90 | "task": f"Superconductor-{oracle}-v0", 91 | "task_kwargs": {"relabel": False, "dataset_kwargs": { 92 | "max_samples": 5000, 93 | "distribution": tune.grid_search([ 94 | "uniform", 95 | "linear", 96 | "quadratic", 97 | "exponential", 98 | "circular" 99 | ]), 100 | "max_percentile": 100, 101 | "min_percentile": 0 102 | }}, 103 | "normalize_ys": True, 104 | "normalize_xs": True, 105 | "model_noise_std": 0.0, 106 | "val_size": 200, 107 | "batch_size": 128, 108 | "epochs": 100, 109 | "activations": [['leaky_relu', 'leaky_relu']] * 5, 110 | "hidden_size": 2048, 111 | "initial_max_std": 0.2, 112 | "initial_min_std": 0.1, 113 | "forward_model_lr": 0.0003, 114 | "aggregation_method": 'min', 115 | "solver_samples": 512, 116 | "solver_lr": 0.01, 117 | "solver_steps": 200}, 118 | num_samples=num_samples, 119 | local_dir=local_dir, 120 | resources_per_trial={'cpu': cpus // num_parallel, 121 | 'gpu': gpus / num_parallel - 0.01}) 122 | 123 | 124 | @cli.command() 125 | @click.option('--local-dir', type=str, default='gradient-ascent-gfp') 126 | @click.option('--cpus', type=int, default=24) 127 | @click.option('--gpus', type=int, default=1) 128 | @click.option('--num-parallel', type=int, default=1) 129 | @click.option('--num-samples', type=int, default=1) 130 | @click.option('--oracle', type=str, default="Transformer") 131 | def gfp(local_dir, cpus, gpus, num_parallel, num_samples, oracle): 132 | """Evaluate Naive Gradient Ascent on GFP-Transformer-v0 133 | """ 134 | 135 | # Final Version 136 | 137 | from design_baselines.gradient_ascent import gradient_ascent 138 | ray.init(num_cpus=cpus, 139 | num_gpus=gpus, 140 | include_dashboard=False, 141 | _temp_dir=os.path.expanduser('~/tmp')) 142 | tune.run(gradient_ascent, config={ 143 | "logging_dir": "data", 144 | "task": f"GFP-{oracle}-v0", 145 | "task_kwargs": {"relabel": False, "dataset_kwargs": { 146 | "max_samples": 5000, 147 | "distribution": tune.grid_search([ 148 | "uniform", 149 | "linear", 150 | "quadratic", 151 | "exponential", 152 | "circular" 153 | ]), 154 | "max_percentile": 100, 155 | "min_percentile": 0 156 | }}, 157 | "normalize_ys": True, 158 | "normalize_xs": True, 159 | "model_noise_std": 0.0, 160 | "val_size": 200, 161 | "use_vae": False, 162 | "vae_beta": 0.01, 163 | "vae_epochs": 50, 164 | "vae_batch_size": 128, 165 | "vae_hidden_size": 64, 166 | "vae_latent_size": 256, 167 | "vae_activation": "relu", 168 | "vae_kernel_size": 3, 169 | "vae_num_blocks": 5, 170 | "vae_lr": 0.0003, 171 | "batch_size": 128, 172 | "epochs": 100, 173 | "activations": [['leaky_relu', 'leaky_relu']] * 5, 174 | "hidden_size": 2048, 175 | "initial_max_std": 0.2, 176 | "initial_min_std": 0.1, 177 | "forward_model_lr": 0.0003, 178 | "aggregation_method": 'min', 179 | "solver_samples": 512, 180 | "solver_lr": 0.01, 181 | "solver_steps": 200}, 182 | num_samples=num_samples, 183 | local_dir=local_dir, 184 | resources_per_trial={'cpu': cpus // num_parallel, 185 | 'gpu': gpus / num_parallel - 0.01}) 186 | 187 | 188 | @cli.command() 189 | @click.option('--local-dir', type=str, default='gradient-ascent-utr') 190 | @click.option('--cpus', type=int, default=24) 191 | @click.option('--gpus', type=int, default=1) 192 | @click.option('--num-parallel', type=int, default=1) 193 | @click.option('--num-samples', type=int, default=1) 194 | def utr(local_dir, cpus, gpus, num_parallel, num_samples): 195 | """Evaluate Naive Gradient Ascent on UTR-ResNet-v0 196 | """ 197 | 198 | # Final Version 199 | 200 | from design_baselines.gradient_ascent import gradient_ascent 201 | ray.init(num_cpus=cpus, 202 | num_gpus=gpus, 203 | include_dashboard=False, 204 | _temp_dir=os.path.expanduser('~/tmp')) 205 | tune.run(gradient_ascent, config={ 206 | "logging_dir": "data", 207 | "task": "UTR-ResNet-v0", 208 | "task_kwargs": {"relabel": False, "dataset_kwargs": { 209 | "max_samples": 20000, 210 | "distribution": tune.grid_search([ 211 | "uniform", 212 | "linear", 213 | "quadratic", 214 | "exponential", 215 | "circular" 216 | ]), 217 | "max_percentile": 100, 218 | "min_percentile": 0 219 | }}, 220 | "normalize_ys": True, 221 | "normalize_xs": True, 222 | "model_noise_std": 0.0, 223 | "val_size": 200, 224 | "use_vae": False, 225 | "vae_beta": 0.01, 226 | "vae_epochs": 50, 227 | "vae_batch_size": 128, 228 | "vae_hidden_size": 64, 229 | "vae_latent_size": 256, 230 | "vae_activation": "relu", 231 | "vae_kernel_size": 3, 232 | "vae_num_blocks": 5, 233 | "vae_lr": 0.0003, 234 | "batch_size": 128, 235 | "epochs": 100, 236 | "activations": [['leaky_relu', 'leaky_relu']] * 5, 237 | "hidden_size": 2048, 238 | "initial_max_std": 0.2, 239 | "initial_min_std": 0.1, 240 | "forward_model_lr": 0.0003, 241 | "aggregation_method": 'min', 242 | "solver_samples": 512, 243 | "solver_lr": 0.01, 244 | "solver_steps": 200}, 245 | num_samples=num_samples, 246 | local_dir=local_dir, 247 | resources_per_trial={'cpu': cpus // num_parallel, 248 | 'gpu': gpus / num_parallel - 0.01}) 249 | -------------------------------------------------------------------------------- /design_baselines/logger.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | import tensorflow_probability as tfp 3 | 4 | 5 | class Logger(object): 6 | 7 | def __init__(self, 8 | logging_dir): 9 | """Creates a logging interface to a tensorboard file for 10 | visualizing in the tensorboard web interface; note that 11 | mean, max, min, and std are recorded 12 | 13 | Arguments: 14 | 15 | logging_dir: str 16 | the path on the disk to save records to 17 | """ 18 | 19 | tf.io.gfile.makedirs(logging_dir) 20 | self.writer = tf.summary.create_file_writer(logging_dir) 21 | 22 | def record(self, 23 | key, 24 | value, 25 | step, 26 | percentile=False): 27 | """Log statistics about training data to tensorboard 28 | log files for visualization later 29 | 30 | Arguments: 31 | 32 | key: str 33 | the string name to use when logging data in tensorboard 34 | that determines groupings in the web interface 35 | value: tf.tensor 36 | the tensor of values to record statistics about 37 | typically is multi dimensional 38 | step: int 39 | the total number of environment steps collected so far 40 | typically on intervals of 10000 41 | """ 42 | 43 | step = tf.cast(tf.convert_to_tensor(step), tf.int64) 44 | with self.writer.as_default(): 45 | 46 | if tf.equal(tf.size(value), 1): 47 | 48 | # log one statistic of the incoming tensors 49 | tf.summary.scalar(key, tf.reshape(value, []), step=step) 50 | 51 | elif percentile: 52 | 53 | # log several statistics of the incoming tensors 54 | tf.summary.scalar(key + '/100th', 55 | tfp.stats.percentile(value, 100.0), 56 | step=step) 57 | tf.summary.scalar(key + '/90th', 58 | tfp.stats.percentile(value, 90.0), 59 | step=step) 60 | tf.summary.scalar(key + '/80th', 61 | tfp.stats.percentile(value, 80.0), 62 | step=step) 63 | tf.summary.scalar(key + '/50th', 64 | tfp.stats.percentile(value, 50.0), 65 | step=step) 66 | 67 | else: 68 | 69 | # log several statistics of the incoming tensors 70 | tf.summary.scalar(key + '/max', 71 | tf.math.reduce_max(value), 72 | step=step) 73 | tf.summary.scalar(key + '/mean', 74 | tf.math.reduce_mean(value), 75 | step=step) 76 | tf.summary.scalar(key + '/min', 77 | tf.math.reduce_min(value), 78 | step=step) 79 | tf.summary.scalar(key + '/std', 80 | tf.math.reduce_std(value), 81 | step=step) 82 | -------------------------------------------------------------------------------- /design_baselines/mins/replay_buffer.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | 3 | 4 | class ReplayBuffer(tf.Module): 5 | 6 | def __init__(self, capacity, shape): 7 | """A static graph replay buffer that stores samples collected 8 | from a generator in a GAN 9 | 10 | Args: 11 | 12 | capacity: tf.dtypes.int32 13 | the number of samples that can be in the buffer, maximum 14 | smaller number = the algorithm is more online 15 | shape: tf.dtypes.int32 16 | the shape of the tensors stored in the buffer 17 | this will be a 3-tensor for images 18 | """ 19 | 20 | super(ReplayBuffer, self).__init__() 21 | self.capacity = capacity 22 | self.shape = shape 23 | 24 | # save size statistics for the buffer 25 | self.head = tf.Variable(tf.constant(0)) 26 | self.size = tf.Variable(tf.constant(0)) 27 | self.step = tf.Variable(tf.constant(0)) 28 | 29 | # prepare a storage memory for samples 30 | self.xs = tf.Variable( 31 | tf.zeros([capacity, *shape], tf.dtypes.float32)) 32 | self.ys = tf.Variable( 33 | tf.zeros([capacity, 1], tf.dtypes.float32)) 34 | 35 | # a tensor for indexing into scatter_nd_update 36 | self.idx = tf.concat([tf.tile(tf.reshape( 37 | tf.range(shape_i), 38 | [1] * i + [shape_i] + [1] * (len(shape) - i)), 39 | list(shape[:i]) + [1] + list(shape[i + 1:]) + [1]) 40 | for i, shape_i in enumerate(shape)], axis=len(shape)) 41 | 42 | @tf.function 43 | def insert(self, x, y): 44 | """Insert a single sample collected from the environment into 45 | the replay buffer at the current head position 46 | 47 | Args: 48 | 49 | x: tf.dtypes.float32 50 | a tensor corresponding to a fake sample from the generator 51 | images may be shaped like [height, width, channels] 52 | y: tf.dtypes.float32 53 | a tensor corresponding to a fake scores from the generator 54 | images may be shaped like [1] 55 | """ 56 | 57 | # create scatter indices for one sample 58 | loc_x = tf.reshape(self.head, [1] * (len(self.shape) + 1)) 59 | loc_x = tf.broadcast_to(loc_x, tf.shape(self.idx))[..., 0:1] 60 | loc_x = tf.concat([loc_x, self.idx], axis=len(self.shape)) 61 | loc_y = tf.concat([tf.reshape(self.head, [1, 1]), [[0]]], axis=1) 62 | 63 | # insert samples at the position of the current head 64 | self.xs.assign(tf.tensor_scatter_nd_update( 65 | self.xs, loc_x, tf.cast(x, tf.dtypes.float32))) 66 | self.ys.assign(tf.tensor_scatter_nd_update( 67 | self.ys, loc_y, tf.cast(y, tf.dtypes.float32))) 68 | 69 | # increment the size statistics of the buffer 70 | self.head.assign( 71 | tf.math.floormod(self.head + 1, self.capacity)) 72 | self.size.assign( 73 | tf.minimum(self.size + 1, self.capacity)) 74 | self.step.assign( 75 | self.step + 1) 76 | 77 | @tf.function 78 | def insert_many(self, xs, ys): 79 | """Insert a single sample collected from the environment into 80 | the replay buffer at the current head position 81 | 82 | Args: 83 | 84 | xs: tf.dtypes.float32 85 | a tensor corresponding to many fake samples from the generator 86 | images may be shaped like [batch, height, width, channels] 87 | ys: tf.dtypes.float32 88 | a tensor corresponding to a fake scores from the generator 89 | images may be shaped like [batch, 1] 90 | """ 91 | 92 | for i in tf.range(tf.shape(xs)[0]): 93 | self.insert(xs[i], ys[i]) 94 | 95 | @tf.function 96 | def sample(self, batch_size): 97 | """Samples a batch of training data from the replay buffer 98 | and returns the batch of data 99 | 100 | Args: 101 | 102 | batch_size: tf.dtypes.int32 103 | a scalar tensor that specifies how many elements to sample 104 | typically smaller than the replay buffer capacity 105 | 106 | Returns: 107 | 108 | xs: tf.dtypes.float32 109 | a tensor corresponding to many fake samples from the generator 110 | images may be shaped like [batch, height, width, channels] 111 | ys: tf.dtypes.float32 112 | a tensor corresponding to a fake scores from the generator 113 | images may be shaped like [batch, 1] 114 | """ 115 | 116 | indices = tf.random.uniform([ 117 | batch_size], maxval=self.size, dtype=tf.dtypes.int32) 118 | return tf.gather(self.xs, indices, axis=0), \ 119 | tf.gather(self.ys, indices, axis=0) 120 | -------------------------------------------------------------------------------- /design_baselines/mins/utils.py: -------------------------------------------------------------------------------- 1 | from tensorflow_probability import distributions as tfpd 2 | import numpy as np 3 | import tensorflow as tf 4 | 5 | 6 | def adaptive_temp_v2(scores_np): 7 | """Calculate an adaptive temperature value based on the 8 | statistics of the scores array 9 | 10 | Args: 11 | 12 | scores_np: np.ndarray 13 | an array that represents the vectorized scores per data point 14 | 15 | Returns: 16 | 17 | temp: np.ndarray 18 | the scalar 90th percentile of scores in the dataset 19 | """ 20 | 21 | inverse_arr = scores_np 22 | max_score = inverse_arr.max() 23 | scores_new = inverse_arr - max_score 24 | quantile_ninety = np.quantile(scores_new, q=0.9) 25 | return np.maximum(np.abs(quantile_ninety), 0.001) 26 | 27 | 28 | def softmax(arr, 29 | temp=1.0): 30 | """Calculate the softmax using numpy by normalizing a vector 31 | to have entries that sum to one 32 | 33 | Args: 34 | 35 | arr: np.ndarray 36 | the array which will be normalized using a tempered softmax 37 | temp: float 38 | a temperature parameter for the softmax 39 | 40 | Returns: 41 | 42 | normalized: np.ndarray 43 | the normalized input array which sums to one 44 | """ 45 | 46 | max_arr = arr.max() 47 | arr_new = arr - max_arr 48 | exp_arr = np.exp(arr_new / temp) 49 | return exp_arr / np.sum(exp_arr) 50 | 51 | 52 | def get_weights(scores, base_temp=None): 53 | """Calculate weights used for training a model inversion 54 | network with a per-sample reweighted objective 55 | 56 | Args: 57 | 58 | scores: np.ndarray 59 | scores which correspond to the value of data points in the dataset 60 | 61 | Returns: 62 | 63 | weights: np.ndarray 64 | an array with the same shape as scores that reweights samples 65 | """ 66 | 67 | scores_np = scores[:, 0] 68 | hist, bin_edges = np.histogram(scores_np, bins=20) 69 | hist = hist / np.sum(hist) 70 | 71 | if base_temp is None: 72 | base_temp = adaptive_temp_v2(scores_np) 73 | softmin_prob = softmax(bin_edges[1:], temp=base_temp) 74 | 75 | provable_dist = softmin_prob * (hist / (hist + 1e-3)) 76 | provable_dist = provable_dist / (np.sum(provable_dist) + 1e-7) 77 | 78 | bin_indices = np.digitize(scores_np, bin_edges[1:]) 79 | hist_prob = hist[np.minimum(bin_indices, 19)] 80 | 81 | weights = provable_dist[ 82 | np.minimum(bin_indices, 19)] / (hist_prob + 1e-7) 83 | weights = np.clip(weights, a_min=0.0, a_max=5.0) 84 | return weights.astype(np.float32)[:, np.newaxis] 85 | 86 | 87 | def get_p_y(scores, base_temp=None): 88 | """Calculate weights used for training a model inversion 89 | network with a per-sample reweighted objective 90 | 91 | Args: 92 | 93 | scores: np.ndarray 94 | scores which correspond to the value of data points in the dataset 95 | 96 | Returns: 97 | 98 | weights: np.ndarray 99 | an array with the same shape as scores that reweights samples 100 | """ 101 | 102 | scores_np = scores[:, 0] 103 | hist, bin_edges = np.histogram(scores_np, bins=20) 104 | hist = hist / np.sum(hist) 105 | 106 | if base_temp is None: 107 | base_temp = adaptive_temp_v2(scores_np) 108 | softmin_prob = softmax(bin_edges[1:], temp=base_temp) 109 | 110 | provable_dist = softmin_prob * (hist / (hist + 1e-3)) 111 | provable_dist = provable_dist / (np.sum(provable_dist) + 1e-7) 112 | 113 | bin_edges = bin_edges.astype(np.float32) 114 | provable_dist = provable_dist.astype(np.float32) 115 | 116 | return provable_dist[:, np.newaxis], \ 117 | bin_edges[:, np.newaxis] 118 | 119 | 120 | @tf.function(experimental_relax_shapes=True) 121 | def get_synthetic_data(x, 122 | y, 123 | exploration_samples=32, 124 | exploration_rate=10.0, 125 | base_temp=None): 126 | """Generate a synthetic dataset of designs x and scores y using the 127 | Randomized Labelling algorithm 128 | 129 | Args: 130 | 131 | x: tf.Tensor 132 | a tensor containing an existing data set of realistic designs 133 | y: tf.Tensor 134 | a tensor containing an existing data set of realistic scores 135 | exploration_samples: int 136 | the number of samples to draw for randomized labelling 137 | exploration_rate: float 138 | the rate of the exponential noise added to y 139 | 140 | Returns: 141 | 142 | tilde_x: tf.Tensor 143 | a tensor containing a data set of synthetic designs 144 | tilde_y: tf.Tensor 145 | a tensor containing a data set of synthetic scores 146 | """ 147 | 148 | def wrapped_py(_y): 149 | return get_p_y(_y, base_temp=base_temp) 150 | 151 | # sample ys based on their importance weight 152 | p_y, bin_edges = tf.numpy_function(wrapped_py, [y], [tf.float32, tf.float32]) 153 | p_y.set_shape([20, 1]) 154 | bin_edges.set_shape([20, 1]) 155 | 156 | # sample ys according to the optimal bins 157 | d = tfpd.Categorical(probs=p_y[:, 0]) 158 | ys_ids = d.sample(sample_shape=(exploration_samples,)) 159 | ys = tf.nn.embedding_lookup(bin_edges, ys_ids) 160 | 161 | # add positive noise to the sampled ys 162 | d = tfpd.Exponential(rate=exploration_rate) 163 | ys = ys + d.sample(sample_shape=tf.shape(ys)) 164 | 165 | # select data points randomly from the data set 166 | d = tfpd.Categorical(logits=tf.zeros([tf.shape(x)[0]])) 167 | xs_ids = d.sample(sample_shape=(exploration_samples,)) 168 | xs = tf.nn.embedding_lookup(x, xs_ids) 169 | 170 | # concatenate newly paired samples with the existing data set 171 | return tf.concat([x, xs], 0), \ 172 | tf.concat([y, ys], 0) 173 | -------------------------------------------------------------------------------- /design_baselines/reinforce/__init__.py: -------------------------------------------------------------------------------- 1 | from design_baselines.data import StaticGraphTask, build_pipeline 2 | from design_baselines.logger import Logger 3 | from design_baselines.reinforce.trainers import Ensemble 4 | from design_baselines.reinforce.nets import ForwardModel 5 | from design_baselines.reinforce.nets import DiscreteMarginal 6 | from design_baselines.reinforce.nets import ContinuousMarginal 7 | from design_baselines.utils import render_video 8 | import tensorflow as tf 9 | import numpy as np 10 | import os 11 | 12 | 13 | def reinforce(config): 14 | """Optimizes over designs x in an offline optimization problem 15 | using the REINFORCE policy gradient method 16 | 17 | Args: 18 | 19 | config: dict 20 | a dictionary of hyper parameters such as the learning rate 21 | """ 22 | 23 | logger = Logger(config['logging_dir']) 24 | task = StaticGraphTask(config['task'], **config['task_kwargs']) 25 | if task.is_discrete: 26 | task.map_to_integers() 27 | 28 | if config['normalize_ys']: 29 | task.map_normalize_y() 30 | if config['normalize_xs']: 31 | task.map_normalize_x() 32 | 33 | x = task.x 34 | y = task.y 35 | 36 | # create the training task and logger 37 | train_data, val_data = build_pipeline( 38 | x=x, y=y, bootstraps=config['bootstraps'], 39 | batch_size=config['ensemble_batch_size'], 40 | val_size=config['val_size']) 41 | 42 | # make several keras neural networks with two hidden layers 43 | forward_models = [ForwardModel( 44 | task, 45 | embedding_size=config['embedding_size'], 46 | hidden_size=config['hidden_size'], 47 | num_layers=config['num_layers'], 48 | initial_max_std=config['initial_max_std'], 49 | initial_min_std=config['initial_min_std']) 50 | for b in range(config['bootstraps'])] 51 | 52 | # create a trainer for a forward model with a conservative objective 53 | ensemble = Ensemble( 54 | forward_models, 55 | forward_model_optim=tf.keras.optimizers.Adam, 56 | forward_model_lr=config['ensemble_lr']) 57 | 58 | # train the model for an additional number of epochs 59 | ensemble.launch(train_data, 60 | val_data, 61 | logger, 62 | config['ensemble_epochs']) 63 | 64 | rl_opt = tf.keras.optimizers.Adam( 65 | learning_rate=config['reinforce_lr']) 66 | 67 | # select the top 1 initial designs from the dataset 68 | indices = tf.math.top_k(y[:, 0], k=config['solver_samples'])[1] 69 | initial_x = tf.gather(x, indices, axis=0) 70 | 71 | if task.is_discrete: 72 | logits = tf.pad(task.to_logits(initial_x), [[0, 0], [0, 0], [1, 0]]) 73 | probs = tf.math.softmax(logits / 1e-5) 74 | logits = tf.math.log(tf.reduce_mean(probs, axis=0)) 75 | sampler = DiscreteMarginal(logits) 76 | 77 | else: 78 | mean = tf.reduce_mean(initial_x, axis=0) 79 | logstd = tf.math.log(tf.ones_like(mean) * config['exploration_std']) 80 | sampler = ContinuousMarginal(mean, logstd) 81 | 82 | for iteration in range(config['iterations']): 83 | 84 | with tf.GradientTape() as tape: 85 | td = sampler.get_distribution() 86 | tx = td.sample(sample_shape=config['reinforce_batch_size']) 87 | if config['optimize_ground_truth']: 88 | ty = task.predict(tx) 89 | else: # use the surrogate model for optimization 90 | ty = ensemble.get_distribution(tx).mean() 91 | 92 | mean_y = tf.reduce_mean(ty) 93 | standard_dev_y = tf.math.reduce_std(ty - mean_y) 94 | log_probs = td.log_prob(tf.stop_gradient(tx)) 95 | loss = tf.reduce_mean(-log_probs[:, tf.newaxis] * 96 | tf.stop_gradient( 97 | (ty - mean_y) / standard_dev_y)) 98 | 99 | print(f"[Iteration {iteration}] " 100 | f"Average Prediction = {tf.reduce_mean(ty)}") 101 | 102 | logger.record("reinforce/prediction", 103 | ty, iteration, percentile=True) 104 | logger.record("reinforce/loss", 105 | loss, iteration, percentile=True) 106 | 107 | grads = tape.gradient( 108 | loss, sampler.trainable_variables) 109 | 110 | rl_opt.apply_gradients(zip( 111 | grads, sampler.trainable_variables)) 112 | 113 | td = sampler.get_distribution() 114 | solution = td.sample(sample_shape=config['solver_samples']) 115 | 116 | # save the current solution to the disk 117 | np.save(os.path.join(config["logging_dir"], 118 | f"solution.npy"), solution.numpy()) 119 | if config["do_evaluation"]: 120 | 121 | # evaluate the found solution and record a video 122 | score = task.predict(solution) 123 | if config['normalize_ys']: 124 | score = task.denormalize_y(score) 125 | logger.record( 126 | "score", score, config['iterations'], percentile=True) 127 | -------------------------------------------------------------------------------- /design_baselines/reinforce/distribution_experiments.py: -------------------------------------------------------------------------------- 1 | from ray import tune 2 | import click 3 | import ray 4 | import os 5 | 6 | 7 | @click.group() 8 | def cli(): 9 | """A group of experiments for training Conservative Score Models 10 | and reproducing our ICLR 2021 results. 11 | """ 12 | 13 | 14 | ############# 15 | 16 | 17 | @cli.command() 18 | @click.option('--local-dir', type=str, default='reinforce-hopper') 19 | @click.option('--cpus', type=int, default=24) 20 | @click.option('--gpus', type=int, default=1) 21 | @click.option('--num-parallel', type=int, default=1) 22 | @click.option('--num-samples', type=int, default=1) 23 | def hopper(local_dir, cpus, gpus, num_parallel, num_samples): 24 | """Evaluate reinforce on HopperController-Exact-v0 25 | """ 26 | 27 | # Final Version 28 | 29 | from design_baselines.reinforce import reinforce 30 | ray.init(num_cpus=cpus, 31 | num_gpus=gpus, 32 | include_dashboard=False, 33 | _temp_dir=os.path.expanduser('~/tmp')) 34 | tune.run(reinforce, config={ 35 | "logging_dir": "data", 36 | "normalize_ys": True, 37 | "normalize_xs": True, 38 | "task": "HopperController-Exact-v0", 39 | "task_kwargs": {"relabel": False, "dataset_kwargs": { 40 | "max_samples": 1000, 41 | "distribution": tune.grid_search([ 42 | "uniform", 43 | "linear", 44 | "quadratic", 45 | "exponential", 46 | "circular" 47 | ]), 48 | "max_percentile": 100, 49 | "min_percentile": 0 50 | }}, 51 | "optimize_ground_truth": False, 52 | "bootstraps": 5, 53 | "val_size": 200, 54 | "ensemble_batch_size": 100, 55 | "embedding_size": 256, 56 | "hidden_size": 256, 57 | "num_layers": 1, 58 | "initial_max_std": 0.2, 59 | "initial_min_std": 0.1, 60 | "ensemble_lr": 0.001, 61 | "ensemble_epochs": 0, 62 | "exploration_std": 0.1, 63 | "reinforce_lr": 0.1, 64 | "reinforce_batch_size": 2048, 65 | "iterations": 200, 66 | "solver_samples": 512}, 67 | num_samples=num_samples, 68 | local_dir=local_dir, 69 | resources_per_trial={'cpu': cpus // num_parallel, 70 | 'gpu': gpus / num_parallel - 0.01}) 71 | 72 | 73 | @cli.command() 74 | @click.option('--local-dir', type=str, default='reinforce-superconductor') 75 | @click.option('--cpus', type=int, default=24) 76 | @click.option('--gpus', type=int, default=1) 77 | @click.option('--num-parallel', type=int, default=1) 78 | @click.option('--num-samples', type=int, default=1) 79 | @click.option('--oracle', type=str, default="RandomForest") 80 | def superconductor(local_dir, cpus, gpus, num_parallel, num_samples, oracle): 81 | """Evaluate reinforce on Superconductor-RandomForest-v0 82 | """ 83 | 84 | # Final Version 85 | 86 | from design_baselines.reinforce import reinforce 87 | ray.init(num_cpus=cpus, 88 | num_gpus=gpus, 89 | include_dashboard=False, 90 | _temp_dir=os.path.expanduser('~/tmp')) 91 | tune.run(reinforce, config={ 92 | "logging_dir": "data", 93 | "normalize_ys": True, 94 | "normalize_xs": True, 95 | "task": f"Superconductor-{oracle}-v0", 96 | "task_kwargs": {"relabel": False, "dataset_kwargs": { 97 | "max_samples": 5000, 98 | "distribution": tune.grid_search([ 99 | "uniform", 100 | "linear", 101 | "quadratic", 102 | "exponential", 103 | "circular" 104 | ]), 105 | "max_percentile": 100, 106 | "min_percentile": 0 107 | }}, 108 | "optimize_ground_truth": False, 109 | "bootstraps": 5, 110 | "val_size": 200, 111 | "ensemble_batch_size": 100, 112 | "embedding_size": 256, 113 | "hidden_size": 256, 114 | "num_layers": 1, 115 | "initial_max_std": 0.2, 116 | "initial_min_std": 0.1, 117 | "ensemble_lr": 0.001, 118 | "ensemble_epochs": 100, 119 | "exploration_std": 0.1, 120 | "reinforce_lr": 0.01, 121 | "reinforce_batch_size": 2048, 122 | "iterations": 200, 123 | "solver_samples": 512}, 124 | num_samples=num_samples, 125 | local_dir=local_dir, 126 | resources_per_trial={'cpu': cpus // num_parallel, 127 | 'gpu': gpus / num_parallel - 0.01}) 128 | 129 | 130 | @cli.command() 131 | @click.option('--local-dir', type=str, default='reinforce-gfp') 132 | @click.option('--cpus', type=int, default=24) 133 | @click.option('--gpus', type=int, default=1) 134 | @click.option('--num-parallel', type=int, default=1) 135 | @click.option('--num-samples', type=int, default=1) 136 | @click.option('--oracle', type=str, default="Transformer") 137 | def gfp(local_dir, cpus, gpus, num_parallel, num_samples, oracle): 138 | """Evaluate reinforce on GFP-Transformer-v0 139 | """ 140 | 141 | # Final Version 142 | 143 | from design_baselines.reinforce import reinforce 144 | ray.init(num_cpus=cpus, 145 | num_gpus=gpus, 146 | include_dashboard=False, 147 | _temp_dir=os.path.expanduser('~/tmp')) 148 | tune.run(reinforce, config={ 149 | "logging_dir": "data", 150 | "normalize_ys": True, 151 | "normalize_xs": False, 152 | "task": f"GFP-{oracle}-v0", 153 | "task_kwargs": {"relabel": False, "dataset_kwargs": { 154 | "max_samples": 5000, 155 | "distribution": tune.grid_search([ 156 | "uniform", 157 | "linear", 158 | "quadratic", 159 | "exponential", 160 | "circular" 161 | ]), 162 | "max_percentile": 100, 163 | "min_percentile": 0 164 | }}, 165 | "optimize_ground_truth": False, 166 | "bootstraps": 5, 167 | "val_size": 200, 168 | "ensemble_batch_size": 100, 169 | "embedding_size": 256, 170 | "hidden_size": 256, 171 | "num_layers": 1, 172 | "initial_max_std": 0.2, 173 | "initial_min_std": 0.1, 174 | "ensemble_lr": 0.001, 175 | "ensemble_epochs": 100, 176 | "reinforce_lr": 0.01, 177 | "reinforce_batch_size": 256, 178 | "iterations": 200, 179 | "solver_samples": 512}, 180 | num_samples=num_samples, 181 | local_dir=local_dir, 182 | resources_per_trial={'cpu': cpus // num_parallel, 183 | 'gpu': gpus / num_parallel - 0.01}) 184 | 185 | 186 | @cli.command() 187 | @click.option('--local-dir', type=str, default='reinforce-utr') 188 | @click.option('--cpus', type=int, default=24) 189 | @click.option('--gpus', type=int, default=1) 190 | @click.option('--num-parallel', type=int, default=1) 191 | @click.option('--num-samples', type=int, default=1) 192 | def utr(local_dir, cpus, gpus, num_parallel, num_samples): 193 | """Evaluate reinforce on UTR-ResNet-v0 194 | """ 195 | 196 | # Final Version 197 | 198 | from design_baselines.reinforce import reinforce 199 | ray.init(num_cpus=cpus, 200 | num_gpus=gpus, 201 | include_dashboard=False, 202 | _temp_dir=os.path.expanduser('~/tmp')) 203 | tune.run(reinforce, config={ 204 | "logging_dir": "data", 205 | "normalize_ys": True, 206 | "normalize_xs": False, 207 | "task": "UTR-ResNet-v0", 208 | "task_kwargs": {"relabel": False, "dataset_kwargs": { 209 | "max_samples": 20000, 210 | "distribution": tune.grid_search([ 211 | "uniform", 212 | "linear", 213 | "quadratic", 214 | "exponential", 215 | "circular" 216 | ]), 217 | "max_percentile": 100, 218 | "min_percentile": 0 219 | }}, 220 | "optimize_ground_truth": False, 221 | "bootstraps": 5, 222 | "val_size": 200, 223 | "ensemble_batch_size": 100, 224 | "embedding_size": 256, 225 | "hidden_size": 256, 226 | "num_layers": 1, 227 | "initial_max_std": 0.2, 228 | "initial_min_std": 0.1, 229 | "ensemble_lr": 0.001, 230 | "ensemble_epochs": 100, 231 | "reinforce_lr": 0.01, 232 | "reinforce_batch_size": 256, 233 | "iterations": 200, 234 | "solver_samples": 512}, 235 | num_samples=num_samples, 236 | local_dir=local_dir, 237 | resources_per_trial={'cpu': cpus // num_parallel, 238 | 'gpu': gpus / num_parallel - 0.01}) 239 | -------------------------------------------------------------------------------- /design_baselines/reinforce/nets.py: -------------------------------------------------------------------------------- 1 | from tensorflow_probability import distributions as tfpd 2 | import tensorflow.keras.layers as tfkl 3 | import tensorflow as tf 4 | import numpy as np 5 | 6 | 7 | class ForwardModel(tf.keras.Sequential): 8 | """A Fully Connected Network with 2 trainable layers""" 9 | 10 | distribution = tfpd.Normal 11 | 12 | def __init__(self, task, embedding_size=50, hidden_size=50, 13 | num_layers=1, initial_max_std=1.5, initial_min_std=0.5): 14 | """Create a fully connected architecture using keras that can process 15 | designs and predict a gaussian distribution over scores 16 | 17 | Args: 18 | 19 | task: StaticGraphTask 20 | a model-based optimization task 21 | embedding_size: int 22 | the size of the embedding matrix for discrete tasks 23 | hidden_size: int 24 | the global hidden size of the neural network 25 | num_layers: int 26 | the number of hidden layers 27 | initial_max_std: float 28 | the starting upper bound of the standard deviation 29 | initial_min_std: float 30 | the starting lower bound of the standard deviation 31 | 32 | """ 33 | 34 | self.max_logstd = tf.Variable(tf.fill([1, 1], np.log( 35 | initial_max_std).astype(np.float32)), trainable=True) 36 | self.min_logstd = tf.Variable(tf.fill([1, 1], np.log( 37 | initial_min_std).astype(np.float32)), trainable=True) 38 | 39 | layers = [] 40 | if task.is_discrete: 41 | layers.append(tfkl.Embedding(task.num_classes, embedding_size, 42 | input_shape=task.input_shape)) 43 | layers.append(tfkl.Flatten(input_shape=task.input_shape) 44 | if len(layers) == 0 else tfkl.Flatten()) 45 | for i in range(num_layers): 46 | layers.extend([tfkl.Dense(hidden_size), tfkl.LeakyReLU()]) 47 | 48 | layers.append(tfkl.Dense(2)) 49 | super(ForwardModel, self).__init__(layers) 50 | 51 | def get_params(self, inputs, **kwargs): 52 | """Return a dictionary of parameters for a particular distribution 53 | family such as the mean and variance of a gaussian 54 | 55 | Args: 56 | 57 | inputs: tf.Tensor 58 | a batch of training inputs shaped like [batch_size, channels] 59 | 60 | Returns: 61 | 62 | parameters: dict 63 | a dictionary that contains 'loc' and 'scale_diag' keys 64 | """ 65 | 66 | prediction = super(ForwardModel, self).__call__(inputs, **kwargs) 67 | mean, logstd = tf.split(prediction, 2, axis=-1) 68 | logstd = self.max_logstd - tf.nn.softplus(self.max_logstd - logstd) 69 | logstd = self.min_logstd + tf.nn.softplus(logstd - self.min_logstd) 70 | return {"loc": mean, "scale": tf.math.exp(logstd)} 71 | 72 | def get_distribution(self, inputs, **kwargs): 73 | """Return a distribution over the outputs of this model, for example 74 | a Multivariate Gaussian Distribution 75 | 76 | Args: 77 | 78 | inputs: tf.Tensor 79 | a batch of training inputs shaped like [batch_size, channels] 80 | 81 | Returns: 82 | 83 | distribution: tfp.distribution.Distribution 84 | a tensorflow probability distribution over outputs of the model 85 | """ 86 | 87 | return self.distribution(**self.get_params(inputs, **kwargs)) 88 | 89 | 90 | class ContinuousMarginal(tf.Module): 91 | 92 | distribution = tfpd.MultivariateNormalDiag 93 | 94 | def __init__(self, 95 | initial_mean, 96 | initial_logstd): 97 | """Create a trainable gaussian distribution with diagonal covariance 98 | 99 | Args: 100 | 101 | input_shape: List[int] 102 | the shape of a single tensor input 103 | """ 104 | 105 | self.mean = tf.Variable(initial_mean, trainable=True) 106 | self.logstd = initial_logstd 107 | super(ContinuousMarginal, self).__init__() 108 | 109 | def get_params(self): 110 | """Return a dictionary of parameters for a particular distribution 111 | family such as the mean and variance of a gaussian 112 | 113 | Returns: 114 | 115 | parameters: dict 116 | a dictionary that contains 'loc' and 'scale_diag' keys 117 | """ 118 | 119 | return {"loc": self.mean, 120 | "scale_diag": tf.math.exp(self.logstd)} 121 | 122 | def get_distribution(self): 123 | """Return a distribution over the outputs of this model, for example 124 | a Multivariate Gaussian Distribution 125 | 126 | Returns: 127 | 128 | distribution: tfp.distribution.Distribution 129 | a tensorflow probability distribution over outputs of the model 130 | """ 131 | 132 | return self.distribution(**self.get_params()) 133 | 134 | 135 | class DiscreteMarginal(tf.Module): 136 | 137 | distribution = tfpd.Categorical 138 | 139 | def __init__(self, 140 | initial_logits): 141 | """Create a trainable one hot categorical distribution 142 | 143 | Args: 144 | 145 | input_shape: List[int] 146 | the shape of a single tensor input 147 | """ 148 | 149 | self.logits = tf.Variable(initial_logits, trainable=True) 150 | super(DiscreteMarginal, self).__init__() 151 | 152 | def get_params(self): 153 | """Return a dictionary of parameters for a particular distribution 154 | family such as the one hot categorical distribution 155 | 156 | Returns: 157 | 158 | parameters: dict 159 | a dictionary that contains 'logits' key 160 | """ 161 | 162 | return {"logits": tf.math.log_softmax(self.logits)} 163 | 164 | def get_distribution(self): 165 | """Return a distribution over the outputs of this model, for example 166 | a one hot categorical distribution 167 | 168 | Returns: 169 | 170 | distribution: tfp.distribution.Distribution 171 | a tensorflow probability distribution over outputs of the model 172 | """ 173 | 174 | return self.distribution(**self.get_params()) 175 | -------------------------------------------------------------------------------- /design_baselines/reinforce/trainers.py: -------------------------------------------------------------------------------- 1 | from design_baselines.utils import spearman 2 | from design_baselines.utils import disc_noise 3 | from design_baselines.utils import cont_noise 4 | from collections import defaultdict 5 | from tensorflow_probability import distributions as tfpd 6 | import tensorflow_probability as tfp 7 | import tensorflow as tf 8 | 9 | 10 | class Ensemble(tf.Module): 11 | 12 | def __init__(self, 13 | forward_models, 14 | forward_model_optim=tf.keras.optimizers.Adam, 15 | forward_model_lr=0.001): 16 | """Build a trainer for an ensemble of probabilistic neural networks 17 | trained on bootstraps of a dataset 18 | 19 | Args: 20 | 21 | oracles: List[tf.keras.Model] 22 | a list of keras model that predict distributions over scores 23 | oracle_optim: __class__ 24 | the optimizer class to use for optimizing the oracle model 25 | oracle__lr: float 26 | the learning rate for the oracle model optimizer 27 | """ 28 | 29 | super().__init__() 30 | self.forward_models = forward_models 31 | self.bootstraps = len(forward_models) 32 | 33 | # create optimizers for each model in the ensemble 34 | self.forward_model_optims = [ 35 | forward_model_optim(learning_rate=forward_model_lr) 36 | for i in range(self.bootstraps)] 37 | 38 | def get_distribution(self, 39 | x, 40 | **kwargs): 41 | """Build the mixture distribution implied by the set of oracles 42 | that are trained in this module 43 | 44 | Args: 45 | 46 | x: tf.Tensor 47 | a batch of training inputs shaped like [batch_size, channels] 48 | 49 | Returns: 50 | 51 | distribution: tfpd.Distribution 52 | the mixture of gaussian distributions implied by the oracles 53 | """ 54 | 55 | # get the distribution parameters for all models 56 | params = defaultdict(list) 57 | for fm in self.forward_models: 58 | for key, val in fm.get_params(x, **kwargs).items(): 59 | params[key].append(val) 60 | 61 | # stack the parameters in a new component axis 62 | for key, val in params.items(): 63 | params[key] = tf.stack(val, axis=-1) 64 | 65 | # build the mixture distribution using the family of component one 66 | weights = tf.fill([self.bootstraps], 1 / self.bootstraps) 67 | return tfpd.MixtureSameFamily(tfpd.Categorical( 68 | probs=weights), self.forward_models[0].distribution(**params)) 69 | 70 | @tf.function(experimental_relax_shapes=True) 71 | def train_step(self, 72 | x, 73 | y, 74 | b): 75 | """Perform a training step of gradient descent on an ensemble 76 | using bootstrap weights for each model in the ensemble 77 | 78 | Args: 79 | 80 | x: tf.Tensor 81 | a batch of training inputs shaped like [batch_size, channels] 82 | y: tf.Tensor 83 | a batch of training labels shaped like [batch_size, 1] 84 | b: tf.Tensor 85 | bootstrap indicators shaped like [batch_size, num_oracles] 86 | 87 | Returns: 88 | 89 | statistics: dict 90 | a dictionary that contains logging information 91 | """ 92 | 93 | statistics = dict() 94 | 95 | for i in range(self.bootstraps): 96 | fm = self.forward_models[i] 97 | fm_optim = self.forward_model_optims[i] 98 | 99 | with tf.GradientTape(persistent=True) as tape: 100 | 101 | # calculate the prediction error and accuracy of the model 102 | d = fm.get_distribution(x, training=True) 103 | nll = -d.log_prob(y)[:, 0] 104 | 105 | # evaluate how correct the rank fo the model predictions are 106 | rank_correlation = spearman(y[:, 0], d.mean()[:, 0]) 107 | 108 | # build the total loss and weight by the bootstrap 109 | total_loss = tf.math.divide_no_nan(tf.reduce_sum( 110 | b[:, i] * nll), tf.reduce_sum(b[:, i])) 111 | 112 | grads = tape.gradient(total_loss, fm.trainable_variables) 113 | fm_optim.apply_gradients(zip(grads, fm.trainable_variables)) 114 | 115 | statistics[f'oracle_{i}/train/nll'] = nll 116 | statistics[f'oracle_{i}/train/rank_corr'] = rank_correlation 117 | 118 | return statistics 119 | 120 | @tf.function(experimental_relax_shapes=True) 121 | def validate_step(self, 122 | x, 123 | y): 124 | """Perform a validation step on an ensemble of models 125 | without using bootstrapping weights 126 | 127 | Args: 128 | 129 | x: tf.Tensor 130 | a batch of validation inputs shaped like [batch_size, channels] 131 | y: tf.Tensor 132 | a batch of validation labels shaped like [batch_size, 1] 133 | 134 | Returns: 135 | 136 | statistics: dict 137 | a dictionary that contains logging information 138 | """ 139 | 140 | statistics = dict() 141 | 142 | for i in range(self.bootstraps): 143 | fm = self.forward_models[i] 144 | 145 | # calculate the prediction error and accuracy of the model 146 | d = fm.get_distribution(x, training=False) 147 | nll = -d.log_prob(y)[:, 0] 148 | 149 | # evaluate how correct the rank fo the model predictions are 150 | rank_correlation = spearman(y[:, 0], d.mean()[:, 0]) 151 | 152 | statistics[f'oracle_{i}/validate/nll'] = nll 153 | statistics[f'oracle_{i}/validate/rank_corr'] = rank_correlation 154 | 155 | return statistics 156 | 157 | def train(self, 158 | dataset): 159 | """Perform training using gradient descent on an ensemble 160 | using bootstrap weights for each model in the ensemble 161 | 162 | Args: 163 | 164 | dataset: tf.data.Dataset 165 | the training dataset already batched and prefetched 166 | 167 | Returns: 168 | 169 | loss_dict: dict 170 | a dictionary mapping names to loss values for logging 171 | """ 172 | 173 | statistics = defaultdict(list) 174 | for x, y, b in dataset: 175 | for name, tensor in self.train_step(x, y, b).items(): 176 | statistics[name].append(tensor) 177 | for name in statistics.keys(): 178 | statistics[name] = tf.concat(statistics[name], axis=0) 179 | return statistics 180 | 181 | def validate(self, 182 | dataset): 183 | """Perform validation on an ensemble of models without 184 | using bootstrapping weights 185 | 186 | Args: 187 | 188 | dataset: tf.data.Dataset 189 | the validation dataset already batched and prefetched 190 | 191 | Returns: 192 | 193 | loss_dict: dict 194 | a dictionary mapping names to loss values for logging 195 | """ 196 | 197 | statistics = defaultdict(list) 198 | for x, y in dataset: 199 | for name, tensor in self.validate_step(x, y).items(): 200 | statistics[name].append(tensor) 201 | for name in statistics.keys(): 202 | statistics[name] = tf.concat(statistics[name], axis=0) 203 | return statistics 204 | 205 | def launch(self, 206 | train_data, 207 | validate_data, 208 | logger, 209 | epochs, 210 | start_epoch=0): 211 | """Launch training and validation for the model for the specified 212 | number of epochs, and log statistics 213 | 214 | Args: 215 | 216 | train_data: tf.data.Dataset 217 | the training dataset already batched and prefetched 218 | validate_data: tf.data.Dataset 219 | the validation dataset already batched and prefetched 220 | logger: Logger 221 | an instance of the logger used for writing to tensor board 222 | epochs: int 223 | the number of epochs through the data sets to take 224 | """ 225 | 226 | for e in range(start_epoch, start_epoch + epochs): 227 | for name, loss in self.train(train_data).items(): 228 | logger.record(name, loss, e) 229 | for name, loss in self.validate(validate_data).items(): 230 | logger.record(name, loss, e) 231 | 232 | def get_saveables(self): 233 | """Collects and returns stateful objects that are serializeable 234 | using the tensorflow checkpoint format 235 | 236 | Returns: 237 | 238 | saveables: dict 239 | a dict containing stateful objects compatible with checkpoints 240 | """ 241 | 242 | saveables = dict() 243 | for i in range(self.bootstraps): 244 | saveables[f'forward_model_{i}'] = self.forward_models[i] 245 | saveables[f'forward_model_optim_{i}'] = self.forward_model_optims[i] 246 | return saveables 247 | -------------------------------------------------------------------------------- /design_baselines/utils.py: -------------------------------------------------------------------------------- 1 | from tensorflow_probability import distributions as tfpd 2 | import tensorflow as tf 3 | import tensorflow_probability as tfp 4 | 5 | 6 | @tf.function(experimental_relax_shapes=True) 7 | def get_rank(x): 8 | return tf.cast(tf.argsort(tf.argsort( 9 | x, axis=-1, direction="ASCENDING"), axis=-1) + 1, x.dtype) 10 | 11 | 12 | @tf.function(experimental_relax_shapes=True) 13 | def spearman(a, b): 14 | """Computes the Spearman Rank-Correlation Coefficient for two 15 | continuous-valued tensors with the same shape 16 | 17 | Args: 18 | 19 | a: tf.Tensor 20 | a tensor of any shape whose last axis represents the axis of which 21 | to rank elements of the tensor from 22 | b: tf.Tensor 23 | a tensor of any shape whose last axis represents the axis of which 24 | to rank elements of the tensor from 25 | 26 | Returns: 27 | 28 | rho: tf.Tensor 29 | a tensor with the same shape as the first N - 1 axes of a and b, and 30 | represents the spearman p between a and b 31 | """ 32 | 33 | x = get_rank(a) 34 | y = get_rank(b) 35 | cov = tfp.stats.covariance( 36 | x, y, sample_axis=-1, keepdims=False, event_axis=None) 37 | sd_x = tfp.stats.stddev( 38 | x, sample_axis=-1, keepdims=True, name=None) 39 | sd_y = tfp.stats.stddev( 40 | y, sample_axis=-1, keepdims=True, name=None) 41 | return cov / (sd_x * sd_y) 42 | 43 | 44 | @tf.function(experimental_relax_shapes=True) 45 | def disc_noise(x, keep=0.9, temp=5.0): 46 | """Add noise to a input that is either a continuous value or a probability 47 | distribution over discrete categorical values 48 | 49 | Args: 50 | 51 | x: tf.Tensor 52 | a tensor that will have noise added to it, such that the resulting 53 | tensor is sound given its definition 54 | keep: float 55 | the amount of probability mass to keep on the element that is activated 56 | teh rest is redistributed evenly to all elements 57 | temp: float 58 | the temperature of teh gumbel distribution that is used to corrupt 59 | the input probabilities x 60 | 61 | Returns: 62 | 63 | noisy_x: tf.Tensor 64 | a tensor that has noise added to it, which has the interpretation of 65 | the original tensor (such as a probability distribution) 66 | """ 67 | 68 | p = tf.ones_like(x) 69 | p = p / tf.reduce_sum(p, axis=-1, keepdims=True) 70 | p = keep * x + (1.0 - keep) * p 71 | return tfpd.RelaxedOneHotCategorical(temp, probs=p).sample() 72 | 73 | 74 | @tf.function(experimental_relax_shapes=True) 75 | def soft_noise(x, keep=0.9, temp=5.0): 76 | """Softens a discrete one-hot distribution so that the maximum entry is 77 | keep and the minimum entries are (1 - keep) 78 | 79 | Args: 80 | 81 | x: tf.Tensor 82 | a tensor that will have noise added to it, such that the resulting 83 | tensor is sound given its definition 84 | keep: float 85 | the amount of probability mass to keep on the element that is 86 | activated; the rest is redistributed to all elements 87 | 88 | Returns: 89 | 90 | smooth_x: tf.Tensor 91 | a smoothed version of x that represents a discrete probability 92 | distribution of categorical variables 93 | """ 94 | 95 | p = tf.ones_like(x) 96 | p = p / tf.reduce_sum(p, axis=-1, keepdims=True) 97 | return keep * x + (1.0 - keep) * p 98 | 99 | 100 | @tf.function(experimental_relax_shapes=True) 101 | def cont_noise(x, noise_std=1.0): 102 | """Add noise to a input that is either a continuous value or a probability 103 | distribution over discrete categorical values 104 | 105 | Args: 106 | 107 | x: tf.Tensor 108 | a tensor that will have noise added to it, such that the resulting 109 | tensor is sound given its definition 110 | noise_std: float 111 | the standard deviation of the gaussian noise that will be added 112 | to the continuous design parameters x 113 | 114 | Returns: 115 | 116 | noisy_x: tf.Tensor 117 | a tensor that has noise added to it, which has the interpretation of 118 | the original tensor (such as a probability distribution) 119 | """ 120 | 121 | return x + noise_std * tf.random.normal(tf.shape(x)) 122 | 123 | 124 | def generate_ensemble(num_layers, *activations): 125 | """Given a set of string names and a number of target layers, generate 126 | a list of ensemble architectures with those activations 127 | 128 | Args: 129 | 130 | num_layers: int 131 | the number of hidden layers in the neural network, and also 132 | the number of activation functions 133 | activations: list of str 134 | a list of strings that indicates the candidates for activation 135 | functions at for every layer 136 | 137 | Returns: 138 | 139 | ensemble: list of list of str 140 | a list of architectures, where an architecture is given by a 141 | list of activation function names 142 | """ 143 | 144 | if num_layers == 0: 145 | return [] 146 | if num_layers == 1: 147 | return [[act] for act in activations] 148 | return [[act, *o] for act in activations 149 | for o in generate_ensemble(num_layers - 1, *activations)] 150 | 151 | 152 | def render_video(config, task, solution): 153 | 154 | if config["task"] == "HopperController-v0": 155 | 156 | import gym 157 | import os 158 | import numpy as np 159 | from skvideo.io import FFmpegWriter 160 | out = FFmpegWriter(os.path.join( 161 | config['logging_dir'], f'vid.mp4')) 162 | 163 | weights = [] 164 | for s in task.wrapped_task.stream_shapes: 165 | weights.append(np.reshape(solution[0:np.prod(s)], s)) 166 | solution = solution[np.prod(s):] 167 | 168 | weights.pop(-1) 169 | 170 | def mlp_policy(h): 171 | h = np.tanh(h @ weights[0] + weights[1]) 172 | h = np.tanh(h @ weights[2] + weights[3]) 173 | return h @ weights[4] + weights[5] 174 | 175 | env = gym.make(task.wrapped_task.env_name) 176 | 177 | for i in range(5): 178 | obs, done = env.reset(), False 179 | path_returns = np.zeros([1], dtype=np.float32) 180 | while not done: 181 | obs, rew, done, info = env.step(mlp_policy(obs)) 182 | path_returns += rew.astype(np.float32) 183 | 184 | out.writeFrame(env.render( 185 | mode='rgb_array', height=500, width=500)) 186 | 187 | out.close() 188 | -------------------------------------------------------------------------------- /environment.yml: -------------------------------------------------------------------------------- 1 | name: design-baselines 2 | channels: 3 | - defaults 4 | - conda-forge 5 | dependencies: 6 | - python>=3.7,<3.8 7 | - pip>=20.0 8 | - conda>=4.7 9 | - cudatoolkit=10.1 10 | - cudnn=7.6 11 | - rdkit 12 | - pip: 13 | - -r ./requirements.txt 14 | -------------------------------------------------------------------------------- /fine-tuning/coms-chembl.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | NUM_TRIALS_PER_GPU=2 3 | IFS=, read -ra DEVICES <<< "$CUDA_VISIBLE_DEVICES" 4 | for PARTICLE_LR in 0.1 0.5 1.0 2.0 5.0; do 5 | for OE_LIMIT in 0.1 0.5 1.0 2.0 5.0; do 6 | for DEVICE in "${DEVICES[@]}"; do 7 | for TRIAL in $(seq $NUM_TRIALS_PER_GPU); do 8 | CUDA_VISIBLE_DEVICES=$DEVICE coms \ 9 | --logging-dir ./coms-chembl/coms-chembl-$PARTICLE_LR-$OE_LIMIT/COMs-ChEMBL_MCHC_CHEMBL3885882_MorganFingerprint-RandomForest-v0-$DEVICE-$TRIAL-$RANDOM \ 10 | --task ChEMBL_MCHC_CHEMBL3885882_MorganFingerprint-RandomForest-v0 \ 11 | --no-task-relabel \ 12 | --normalize-ys \ 13 | --normalize-xs \ 14 | --not-in-latent-space \ 15 | --vae-hidden-size 64 \ 16 | --vae-latent-size 256 \ 17 | --vae-activation relu \ 18 | --vae-kernel-size 3 \ 19 | --vae-num-blocks 4 \ 20 | --vae-lr 0.0003 \ 21 | --vae-beta 1.0 \ 22 | --vae-batch-size 128 \ 23 | --vae-val-size 500 \ 24 | --vae-epochs 10 \ 25 | --particle-lr $PARTICLE_LR \ 26 | --particle-train-gradient-steps 50 \ 27 | --particle-evaluate-gradient-steps 50 \ 28 | --particle-entropy-coefficient 0.0 \ 29 | --forward-model-activations relu \ 30 | --forward-model-activations relu \ 31 | --forward-model-hidden-size 2048 \ 32 | --no-forward-model-final-tanh \ 33 | --forward-model-lr 0.0003 \ 34 | --forward-model-alpha 0.1 \ 35 | --forward-model-alpha-lr 0.01 \ 36 | --forward-model-overestimation-limit $OE_LIMIT \ 37 | --forward-model-noise-std 0.0 \ 38 | --forward-model-batch-size 128 \ 39 | --forward-model-val-size 100 \ 40 | --forward-model-epochs 50 \ 41 | --evaluation-samples 128 \ 42 | --fast & done; done; wait 43 | done 44 | done 45 | -------------------------------------------------------------------------------- /fine-tuning/coms-nas.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | NUM_TRIALS_PER_GPU=2 3 | IFS=, read -ra DEVICES <<< "$CUDA_VISIBLE_DEVICES" 4 | for PARTICLE_LR in 0.1 0.5 1.0 2.0 5.0; do 5 | for OE_LIMIT in 0.1 0.5 1.0 2.0 5.0; do 6 | for DEVICE in "${DEVICES[@]}"; do 7 | for TRIAL in $(seq $NUM_TRIALS_PER_GPU); do 8 | CUDA_VISIBLE_DEVICES=$DEVICE coms \ 9 | --logging-dir ./coms-nas/coms-nas-$PARTICLE_LR-$OE_LIMIT/COMs-CIFARNAS-Exact-v0-$DEVICE-$TRIAL-$RANDOM \ 10 | --task CIFARNAS-Exact-v0 \ 11 | --no-task-relabel \ 12 | --normalize-ys \ 13 | --normalize-xs \ 14 | --not-in-latent-space \ 15 | --vae-hidden-size 64 \ 16 | --vae-latent-size 256 \ 17 | --vae-activation relu \ 18 | --vae-kernel-size 3 \ 19 | --vae-num-blocks 4 \ 20 | --vae-lr 0.0003 \ 21 | --vae-beta 1.0 \ 22 | --vae-batch-size 128 \ 23 | --vae-val-size 500 \ 24 | --vae-epochs 10 \ 25 | --particle-lr $PARTICLE_LR \ 26 | --particle-train-gradient-steps 50 \ 27 | --particle-evaluate-gradient-steps 50 \ 28 | --particle-entropy-coefficient 0.0 \ 29 | --forward-model-activations relu \ 30 | --forward-model-activations relu \ 31 | --forward-model-hidden-size 2048 \ 32 | --no-forward-model-final-tanh \ 33 | --forward-model-lr 0.0003 \ 34 | --forward-model-alpha 0.1 \ 35 | --forward-model-alpha-lr 0.01 \ 36 | --forward-model-overestimation-limit $OE_LIMIT \ 37 | --forward-model-noise-std 0.0 \ 38 | --forward-model-batch-size 128 \ 39 | --forward-model-val-size 100 \ 40 | --forward-model-epochs 50 \ 41 | --evaluation-samples 128 \ 42 | --fast & done; done; wait 43 | done 44 | done 45 | -------------------------------------------------------------------------------- /fine-tuning/coms-tf-bind-10.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | NUM_TRIALS_PER_GPU=2 3 | IFS=, read -ra DEVICES <<< "$CUDA_VISIBLE_DEVICES" 4 | for PARTICLE_LR in 0.1 0.5 1.0 2.0 5.0; do 5 | for OE_LIMIT in 0.1 0.5 1.0 2.0 5.0; do 6 | for DEVICE in "${DEVICES[@]}"; do 7 | for TRIAL in $(seq $NUM_TRIALS_PER_GPU); do 8 | CUDA_VISIBLE_DEVICES=$DEVICE coms \ 9 | --logging-dir ./coms-tf-bind-10/coms-tf-bind-10-$PARTICLE_LR-$OE_LIMIT/COMs-TFBind10-Exact-v0-$DEVICE-$TRIAL-$RANDOM \ 10 | --task TFBind10-Exact-v0 \ 11 | --no-task-relabel \ 12 | --task-max-samples 10000 \ 13 | --task-distribution uniform \ 14 | --normalize-ys \ 15 | --normalize-xs \ 16 | --not-in-latent-space \ 17 | --vae-hidden-size 64 \ 18 | --vae-latent-size 256 \ 19 | --vae-activation relu \ 20 | --vae-kernel-size 3 \ 21 | --vae-num-blocks 4 \ 22 | --vae-lr 0.0003 \ 23 | --vae-beta 1.0 \ 24 | --vae-batch-size 128 \ 25 | --vae-val-size 500 \ 26 | --vae-epochs 10 \ 27 | --particle-lr $PARTICLE_LR \ 28 | --particle-train-gradient-steps 50 \ 29 | --particle-evaluate-gradient-steps 50 \ 30 | --particle-entropy-coefficient 0.0 \ 31 | --forward-model-activations relu \ 32 | --forward-model-activations relu \ 33 | --forward-model-hidden-size 2048 \ 34 | --no-forward-model-final-tanh \ 35 | --forward-model-lr 0.0003 \ 36 | --forward-model-alpha 0.1 \ 37 | --forward-model-alpha-lr 0.01 \ 38 | --forward-model-overestimation-limit $OE_LIMIT \ 39 | --forward-model-noise-std 0.0 \ 40 | --forward-model-batch-size 128 \ 41 | --forward-model-val-size 500 \ 42 | --forward-model-epochs 50 \ 43 | --evaluation-samples 128 \ 44 | --fast & done; done; wait 45 | done 46 | done 47 | -------------------------------------------------------------------------------- /fine-tuning/coms-tf-bind-8.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | NUM_TRIALS_PER_GPU=2 3 | IFS=, read -ra DEVICES <<< "$CUDA_VISIBLE_DEVICES" 4 | for PARTICLE_LR in 0.1 0.5 1.0 2.0 5.0; do 5 | for OE_LIMIT in 0.1 0.5 1.0 2.0 5.0; do 6 | for DEVICE in "${DEVICES[@]}"; do 7 | for TRIAL in $(seq $NUM_TRIALS_PER_GPU); do 8 | CUDA_VISIBLE_DEVICES=$DEVICE coms \ 9 | --logging-dir ~/coms-tf-bind-8/coms-tf-bind-8-$PARTICLE_LR-$OE_LIMIT/COMs-TFBind8-Exact-v0-$DEVICE-$TRIAL-$RANDOM \ 10 | --task TFBind8-Exact-v0 \ 11 | --no-task-relabel \ 12 | --normalize-ys \ 13 | --normalize-xs \ 14 | --not-in-latent-space \ 15 | --vae-hidden-size 64 \ 16 | --vae-latent-size 256 \ 17 | --vae-activation relu \ 18 | --vae-kernel-size 3 \ 19 | --vae-num-blocks 4 \ 20 | --vae-lr 0.0003 \ 21 | --vae-beta 1.0 \ 22 | --vae-batch-size 128 \ 23 | --vae-val-size 500 \ 24 | --vae-epochs 10 \ 25 | --particle-lr $PARTICLE_LR \ 26 | --particle-train-gradient-steps 50 \ 27 | --particle-evaluate-gradient-steps 50 \ 28 | --particle-entropy-coefficient 0.0 \ 29 | --forward-model-activations relu \ 30 | --forward-model-activations relu \ 31 | --forward-model-hidden-size 2048 \ 32 | --no-forward-model-final-tanh \ 33 | --forward-model-lr 0.0003 \ 34 | --forward-model-alpha 0.1 \ 35 | --forward-model-alpha-lr 0.01 \ 36 | --forward-model-overestimation-limit $OE_LIMIT \ 37 | --forward-model-noise-std 0.0 \ 38 | --forward-model-batch-size 128 \ 39 | --forward-model-val-size 500 \ 40 | --forward-model-epochs 50 \ 41 | --evaluation-samples 128 \ 42 | --fast & done; done; wait 43 | done 44 | done 45 | -------------------------------------------------------------------------------- /label_with_scores.py: -------------------------------------------------------------------------------- 1 | import pickle as pkl 2 | import numpy as np 3 | import os 4 | import glob 5 | import design_bench as db 6 | 7 | 8 | if __name__ == "__main__": 9 | 10 | task = db.make("CIFARNAS-Exact-v0") 11 | task.map_to_logits() 12 | task.map_normalize_x() 13 | 14 | with open('final_output.pkl', 'rb') as f: 15 | final_output = pkl.load(f) 16 | 17 | baselines = [ 18 | "autofocused-cbas", 19 | "cbas", 20 | "bo-qei", 21 | "cma-es", 22 | "gradient-ascent", 23 | "gradient-ascent-min-ensemble", 24 | "gradient-ascent-mean-ensemble", 25 | "mins", 26 | "reinforce" 27 | ] 28 | 29 | baseline_to_iteration = { 30 | "autofocused-cbas": 20, 31 | "cbas": 20, 32 | "bo-qei": 10, 33 | "cma-es": 0, 34 | "gradient-ascent": 200, 35 | "gradient-ascent-min-ensemble": 200, 36 | "gradient-ascent-mean-ensemble": 200, 37 | "mins": 0, 38 | "reinforce": 200 39 | } 40 | 41 | baseline_to_scores = { 42 | "autofocused-cbas": [], 43 | "cbas": [], 44 | "bo-qei": [], 45 | "cma-es": [], 46 | "gradient-ascent": [], 47 | "gradient-ascent-min-ensemble": [], 48 | "gradient-ascent-mean-ensemble": [], 49 | "mins": [], 50 | "reinforce": [] 51 | } 52 | 53 | min_num_samples = 128 54 | 55 | for baseline in baselines: 56 | 57 | solutions = glob.glob(f"/home/btrabucco/final-results" 58 | f"/{baseline}-nas/*/*/data/solution.npy") 59 | 60 | for soln in solutions: 61 | 62 | data = np.load(soln) 63 | if data.shape == (128, 64, 4): 64 | data = task.to_integers(task.denormalize_x(data)) 65 | 66 | data = ["-".join(row.tolist()) for row in data] 67 | 68 | baseline_to_scores[baseline].append([ 69 | final_output[row] for row in data if row in final_output]) 70 | 71 | print(baseline, ':', baseline_to_scores[baseline]) 72 | 73 | -------------------------------------------------------------------------------- /process_nas_results.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | import numpy as np 3 | import argparse 4 | import os 5 | import random 6 | import string 7 | import pickle as pkl 8 | import glob 9 | import design_bench as db 10 | from datetime import datetime 11 | 12 | 13 | if __name__ == "__main__": 14 | 15 | baseline_to_iteration = { 16 | "autofocused-cbas": 20, 17 | "cbas": 20, 18 | "bo-qei": 10, 19 | "cma-es": 0, 20 | "gradient-ascent": 200, 21 | "gradient-ascent-min-ensemble": 200, 22 | "gradient-ascent-mean-ensemble": 200, 23 | "mins": 0, 24 | "reinforce": 200 25 | } 26 | 27 | parser = argparse.ArgumentParser() 28 | parser.add_argument("--input-dir", type=str, default="/home/btrabucco/mbo-results") 29 | parser.add_argument("--input-pkl", type=str, default="final_output.pkl") 30 | parser.add_argument("--out", type=str, default="/home/btrabucco/final-results") 31 | args = parser.parse_args() 32 | 33 | now = datetime.now().strftime("%Y-%m-%d_%H-%M-%S") 34 | digits_and_letters = list(string.digits) + list(string.ascii_letters) 35 | 36 | task = db.make("CIFARNAS-Exact-v0") 37 | task.map_to_logits() 38 | task.map_normalize_x() 39 | 40 | with open(args.input_pkl, "rb") as f: 41 | input_dict = pkl.load(f) 42 | 43 | for baseline, step in baseline_to_iteration.items(): 44 | 45 | solution_files = glob.glob(os.path.join( 46 | args.input_dir, f"{baseline}-nas/*/*/data/solution.npy")) 47 | 48 | step = tf.cast(tf.convert_to_tensor( 49 | baseline_to_iteration[baseline]), tf.int64) 50 | 51 | for solution_file in solution_files: 52 | solution = np.load(solution_file) 53 | 54 | if solution.shape == (128, 64, 4): 55 | solution = task.to_integers(task.denormalize_x(solution)) 56 | 57 | perf = [input_dict["-".join([str(xi) for xi in x])] 58 | for x in solution.tolist()] 59 | 60 | perf_50 = np.percentile(perf, 50) 61 | perf_80 = np.percentile(perf, 80) 62 | perf_90 = np.percentile(perf, 90) 63 | perf_100 = np.percentile(perf, 100) 64 | 65 | print(baseline, perf_50, perf_80, perf_90, perf_100) 66 | 67 | sweep_folder_name = os.path.basename(os.path.dirname(os.path.dirname(os.path.dirname(solution_file)))) 68 | seed_folder_name = os.path.basename(os.path.dirname(os.path.dirname(solution_file))) 69 | 70 | output_dir = os.path.join( 71 | args.out, 72 | f"{baseline}-cifar-nas/" 73 | f"{sweep_folder_name}/{seed_folder_name}/data/") 74 | 75 | step = tf.cast(tf.convert_to_tensor( 76 | baseline_to_iteration[baseline]), tf.int64) 77 | 78 | tf.io.gfile.makedirs(output_dir) 79 | writer = tf.summary.create_file_writer(output_dir) 80 | 81 | with writer.as_default(): 82 | 83 | tf.summary.scalar( 84 | f'score/50th', perf_50, step=step) 85 | tf.summary.scalar( 86 | f'score/80th', perf_80, step=step) 87 | tf.summary.scalar( 88 | f'score/90th', perf_90, step=step) 89 | tf.summary.scalar( 90 | f'score/100th', perf_100, step=step) 91 | -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | absl-py==0.12.0 2 | aiohttp==3.6.2 3 | aiohttp-cors==0.7.0 4 | aioredis==1.3.1 5 | argon2-cffi==20.1.0 6 | astunparse==1.6.3 7 | async-timeout==3.0.1 8 | attrs==19.3.0 9 | backcall==0.2.0 10 | beautifulsoup4==4.9.1 11 | biopython==1.78 12 | bleach==3.1.5 13 | blessings==1.7 14 | boto3==1.16.19 15 | botocore==1.19.19 16 | botorch==0.3.3 17 | brotlipy==0.7.0 18 | cachetools==4.1.1 19 | certifi==2020.6.20 20 | chardet==3.0.4 21 | click==7.1.2 22 | cloudpickle==1.3.0 23 | cma==3.0.3 24 | colorama==0.4.3 25 | colorful==0.5.4 26 | cycler==0.10.0 27 | Cython==0.29.21 28 | decorator==4.4.2 29 | deepchem==2.5.0 30 | defusedxml==0.6.0 31 | dm-tree==0.1.5 32 | docker==4.3.0 33 | entrypoints==0.3 34 | fasteners==0.16 35 | filelock==3.0.12 36 | flatbuffers==1.12 37 | future==0.18.2 38 | gast==0.3.3 39 | glfw==1.12.0 40 | google==3.0.0 41 | google-api-core==1.22.2 42 | google-auth==1.21.3 43 | google-auth-oauthlib==0.4.1 44 | google-pasta==0.2.0 45 | googleapis-common-protos==1.52.0 46 | gpustat==0.6.0 47 | gpytorch==1.3.0 48 | grpcio==1.32.0 49 | gym==0.12.5 50 | h5py==2.10.0 51 | hiredis==1.1.0 52 | huggingface-hub==0.0.8 53 | imageio==2.9.0 54 | importlib-metadata==1.7.0 55 | ipdb==0.13.4 56 | ipykernel==5.3.4 57 | ipython==7.17.0 58 | ipython-genutils==0.2.0 59 | ipywidgets==7.5.1 60 | jedi==0.17.2 61 | Jinja2==2.11.2 62 | jmespath==0.10.0 63 | joblib==1.0.0 64 | jsonschema==3.2.0 65 | jupyter==1.0.0 66 | jupyter-client==6.1.6 67 | jupyter-console==6.1.0 68 | jupyter-core==4.6.3 69 | Keras-Preprocessing==1.1.2 70 | kiwisolver==1.2.0 71 | lmdb==1.0.0 72 | lockfile==0.12.2 73 | Markdown==3.2.2 74 | MarkupSafe==1.1.1 75 | matplotlib==3.3.0 76 | mistune==0.8.4 77 | msgpack==1.0.0 78 | mujoco-py==2.0.2.3 79 | multidict==4.7.6 80 | nbconvert==5.6.1 81 | nbformat==5.0.7 82 | notebook==6.1.1 83 | numpy==1.18.5 84 | nvidia-ml-py3==7.352.0 85 | oauthlib==3.1.0 86 | opencensus==0.7.10 87 | opencensus-context==0.1.1 88 | opt-einsum==3.3.0 89 | packaging==20.4 90 | pandas==1.0.5 91 | pandocfilters==1.4.2 92 | parso==0.7.1 93 | pexpect==4.8.0 94 | pickleshare==0.7.5 95 | Pillow==7.2.0 96 | prometheus-client==0.8.0 97 | prompt-toolkit==3.0.6 98 | protobuf==3.17.1 99 | psutil==5.7.2 100 | ptyprocess==0.6.0 101 | py-spy==0.3.3 102 | pyasn1==0.4.8 103 | pyasn1-modules==0.2.8 104 | pycosat==0.6.3 105 | pyglet==1.5.0 106 | Pygments==2.6.1 107 | pyparsing==2.4.7 108 | pyrsistent==0.16.0 109 | python-dateutil==2.8.1 110 | pytz==2020.1 111 | PyYAML==5.3.1 112 | pyzmq==19.0.2 113 | qtconsole==4.7.5 114 | QtPy==1.9.0 115 | ray==1.3.0 116 | redis==3.5.3 117 | regex==2021.4.4 118 | requests-oauthlib==1.3.0 119 | robel==0.1.2 120 | rsa==4.6 121 | ruamel.yaml==0.15.87 122 | s3transfer==0.3.3 123 | sacremoses==0.0.45 124 | scikit-learn==0.23.1 125 | scikit-video==1.1.11 126 | scipy==1.6.0 127 | seaborn==0.11.0 128 | Send2Trash==1.5.0 129 | soupsieve==2.0.1 130 | tabulate==0.8.7 131 | tensorboard==2.5.0 132 | tensorboard-data-server==0.6.1 133 | tensorboard-plugin-wit==1.7.0 134 | tensorboardX==2.1 135 | tensorflow==2.3.2 136 | tensorflow-estimator==2.3.0 137 | tensorflow-probability==0.11.0 138 | termcolor==1.1.0 139 | terminado==0.8.3 140 | testpath==0.4.4 141 | threadpoolctl==2.1.0 142 | tokenizers==0.9.3 143 | torch==1.7.1 144 | torchvision==0.8.2 145 | tornado==6.0.4 146 | tqdm==4.56.0 147 | traitlets==4.3.3 148 | transformers==3.5.1 149 | transforms3d==0.3.1 150 | typing-extensions==3.7.4.3 151 | wcwidth==0.2.5 152 | webencodings==0.5.1 153 | websocket-client==0.57.0 154 | Werkzeug==1.0.1 155 | widgetsnbextension==3.5.1 156 | wrapt==1.12.1 157 | yarl==1.4.2 158 | zipp==3.1.0 159 | design-bench[all]==2.0.20 160 | morphing-agents==1.5.1 161 | -e . 162 | -------------------------------------------------------------------------------- /scripts/coms-ant.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | NUM_TRIALS_PER_GPU=2 3 | IFS=, read -ra DEVICES <<< "$CUDA_VISIBLE_DEVICES" 4 | for OE_LIMIT in 0.4 0.5 0.6 0.8 1.0; do 5 | for DEVICE in "${DEVICES[@]}"; do 6 | for TRIAL in $(seq $NUM_TRIALS_PER_GPU); do 7 | CUDA_VISIBLE_DEVICES=$DEVICE coms \ 8 | --logging-dir ~/coms-ant/coms-ant-$OE_LIMIT/COMs-AntMorphology-Exact-v0-$DEVICE-$TRIAL-$RANDOM \ 9 | --task AntMorphology-Exact-v0 \ 10 | --no-task-relabel \ 11 | --normalize-ys \ 12 | --normalize-xs \ 13 | --not-in-latent-space \ 14 | --vae-hidden-size 64 \ 15 | --vae-latent-size 256 \ 16 | --vae-activation relu \ 17 | --vae-kernel-size 3 \ 18 | --vae-num-blocks 4 \ 19 | --vae-lr 0.0003 \ 20 | --vae-beta 1.0 \ 21 | --vae-batch-size 128 \ 22 | --vae-val-size 500 \ 23 | --vae-epochs 10 \ 24 | --particle-lr 0.05 \ 25 | --particle-train-gradient-steps 50 \ 26 | --particle-evaluate-gradient-steps 50 \ 27 | --particle-entropy-coefficient 0.0 \ 28 | --forward-model-activations relu \ 29 | --forward-model-activations relu \ 30 | --forward-model-hidden-size 2048 \ 31 | --no-forward-model-final-tanh \ 32 | --forward-model-lr 0.0003 \ 33 | --forward-model-alpha 0.1 \ 34 | --forward-model-alpha-lr 0.01 \ 35 | --forward-model-overestimation-limit $OE_LIMIT \ 36 | --forward-model-noise-std 0.0 \ 37 | --forward-model-batch-size 128 \ 38 | --forward-model-val-size 500 \ 39 | --forward-model-epochs 50 \ 40 | --evaluation-samples 128 \ 41 | --fast & done; done; wait 42 | done 43 | -------------------------------------------------------------------------------- /scripts/coms-chembl.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | . /packages/anaconda3/etc/profile.d/conda.shi 3 | conda activate design-baselines 4 | NUM_TRIALS_PER_GPU=2 5 | IFS=, read -ra DEVICES <<< "$CUDA_VISIBLE_DEVICES" 6 | for OE_LIMIT in 2.0; do 7 | for DEVICE in "${DEVICES[@]}"; do 8 | for TRIAL in $(seq $NUM_TRIALS_PER_GPU); do 9 | CUDA_VISIBLE_DEVICES=$DEVICE coms \ 10 | --logging-dir ./coms-chembl/coms-chembl-$OE_LIMIT/COMs-ChEMBL_MCHC_CHEMBL3885882_MorganFingerprint-RandomForest-v0-$DEVICE-$TRIAL-$RANDOM \ 11 | --task ChEMBL_MCHC_CHEMBL3885882_MorganFingerprint-RandomForest-v0 \ 12 | --no-task-relabel \ 13 | --normalize-ys \ 14 | --no-normalize-xs \ 15 | --not-in-latent-space \ 16 | --vae-hidden-size 64 \ 17 | --vae-latent-size 256 \ 18 | --vae-activation relu \ 19 | --vae-kernel-size 3 \ 20 | --vae-num-blocks 4 \ 21 | --vae-lr 0.0003 \ 22 | --vae-beta 1.0 \ 23 | --vae-batch-size 128 \ 24 | --vae-val-size 500 \ 25 | --vae-epochs 10 \ 26 | --particle-lr 2.0 \ 27 | --particle-train-gradient-steps 50 \ 28 | --particle-evaluate-gradient-steps 50 \ 29 | --particle-entropy-coefficient 0.0 \ 30 | --forward-model-activations relu \ 31 | --forward-model-activations relu \ 32 | --forward-model-hidden-size 2048 \ 33 | --no-forward-model-final-tanh \ 34 | --forward-model-lr 0.0003 \ 35 | --forward-model-alpha 0.1 \ 36 | --forward-model-alpha-lr 0.01 \ 37 | --forward-model-overestimation-limit $OE_LIMIT \ 38 | --forward-model-noise-std 0.0 \ 39 | --forward-model-batch-size 128 \ 40 | --forward-model-val-size 100 \ 41 | --forward-model-epochs 50 \ 42 | --evaluation-samples 128 \ 43 | --fast & done; done; wait 44 | done 45 | -------------------------------------------------------------------------------- /scripts/coms-dkitty.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | NUM_TRIALS_PER_GPU=2 3 | IFS=, read -ra DEVICES <<< "$CUDA_VISIBLE_DEVICES" 4 | for OE_LIMIT in 0.4 0.5 0.6 0.8 1.0; do 5 | for DEVICE in "${DEVICES[@]}"; do 6 | for TRIAL in $(seq $NUM_TRIALS_PER_GPU); do 7 | CUDA_VISIBLE_DEVICES=$DEVICE coms \ 8 | --logging-dir ~/coms-dkitty/coms-dkitty-$OE_LIMIT/COMs-DKittyMorphology-Exact-v0-$DEVICE-$TRIAL-$RANDOM \ 9 | --task DKittyMorphology-Exact-v0 \ 10 | --no-task-relabel \ 11 | --normalize-ys \ 12 | --normalize-xs \ 13 | --not-in-latent-space \ 14 | --vae-hidden-size 64 \ 15 | --vae-latent-size 256 \ 16 | --vae-activation relu \ 17 | --vae-kernel-size 3 \ 18 | --vae-num-blocks 4 \ 19 | --vae-lr 0.0003 \ 20 | --vae-beta 1.0 \ 21 | --vae-batch-size 128 \ 22 | --vae-val-size 500 \ 23 | --vae-epochs 10 \ 24 | --particle-lr 0.05 \ 25 | --particle-train-gradient-steps 50 \ 26 | --particle-evaluate-gradient-steps 50 \ 27 | --particle-entropy-coefficient 0.0 \ 28 | --forward-model-activations relu \ 29 | --forward-model-activations relu \ 30 | --forward-model-hidden-size 2048 \ 31 | --no-forward-model-final-tanh \ 32 | --forward-model-lr 0.0003 \ 33 | --forward-model-alpha 0.1 \ 34 | --forward-model-alpha-lr 0.01 \ 35 | --forward-model-overestimation-limit $OE_LIMIT \ 36 | --forward-model-noise-std 0.0 \ 37 | --forward-model-batch-size 128 \ 38 | --forward-model-val-size 500 \ 39 | --forward-model-epochs 50 \ 40 | --evaluation-samples 128 \ 41 | --fast & done; done; wait 42 | done 43 | -------------------------------------------------------------------------------- /scripts/coms-gfp.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | NUM_TRIALS_PER_GPU=2 3 | IFS=, read -ra DEVICES <<< "$CUDA_VISIBLE_DEVICES" 4 | for OE_LIMIT in 2.0; do 5 | for DEVICE in "${DEVICES[@]}"; do 6 | for TRIAL in $(seq $NUM_TRIALS_PER_GPU); do 7 | CUDA_VISIBLE_DEVICES=$DEVICE coms \ 8 | --logging-dir ~/coms-gfp/coms-gfp-$OE_LIMIT/COMs-GFP-Transformer-v0-$DEVICE-$TRIAL-$RANDOM \ 9 | --task GFP-Transformer-v0 \ 10 | --no-task-relabel \ 11 | --normalize-ys \ 12 | --no-normalize-xs \ 13 | --not-in-latent-space \ 14 | --vae-hidden-size 64 \ 15 | --vae-latent-size 256 \ 16 | --vae-activation relu \ 17 | --vae-kernel-size 3 \ 18 | --vae-num-blocks 4 \ 19 | --vae-lr 0.0003 \ 20 | --vae-beta 1.0 \ 21 | --vae-batch-size 128 \ 22 | --vae-val-size 500 \ 23 | --vae-epochs 10 \ 24 | --particle-lr 2.0 \ 25 | --particle-train-gradient-steps 50 \ 26 | --particle-evaluate-gradient-steps 50 \ 27 | --particle-entropy-coefficient 0.0 \ 28 | --forward-model-activations relu \ 29 | --forward-model-activations relu \ 30 | --forward-model-hidden-size 2048 \ 31 | --no-forward-model-final-tanh \ 32 | --forward-model-lr 0.0003 \ 33 | --forward-model-alpha 0.1 \ 34 | --forward-model-alpha-lr 0.01 \ 35 | --forward-model-overestimation-limit $OE_LIMIT \ 36 | --forward-model-noise-std 0.0 \ 37 | --forward-model-batch-size 128 \ 38 | --forward-model-val-size 500 \ 39 | --forward-model-epochs 50 \ 40 | --evaluation-samples 128 \ 41 | --fast & done; done; wait 42 | done 43 | -------------------------------------------------------------------------------- /scripts/coms-hopper-ablate-conservatism.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | NUM_TRIALS_PER_GPU=2 3 | IFS=, read -ra DEVICES <<< "$CUDA_VISIBLE_DEVICES" 4 | 5 | for OE_LIMIT in 0.1 0.2 0.5 1.0 2.0; do 6 | for DEVICE in "${DEVICES[@]}"; do 7 | for TRIAL in $(seq $NUM_TRIALS_PER_GPU); do 8 | CUDA_VISIBLE_DEVICES=$DEVICE coms \ 9 | --logging-dir ~/coms-hopper-ablate/coms-hopper-$OE_LIMIT/COMs-HopperController-Exact-v0-$DEVICE-$TRIAL-$RANDOM \ 10 | --task HopperController-Exact-v0 \ 11 | --no-task-relabel \ 12 | --normalize-ys \ 13 | --normalize-xs \ 14 | --not-in-latent-space \ 15 | --particle-lr 0.05 \ 16 | --particle-train-gradient-steps 50 \ 17 | --particle-evaluate-gradient-steps 50 \ 18 | --particle-entropy-coefficient 0.0 \ 19 | --forward-model-activations relu \ 20 | --forward-model-activations relu \ 21 | --forward-model-hidden-size 2048 \ 22 | --no-forward-model-final-tanh \ 23 | --forward-model-lr 0.0003 \ 24 | --forward-model-alpha 0.1 \ 25 | --forward-model-alpha-lr 0.01 \ 26 | --forward-model-overestimation-limit $OE_LIMIT \ 27 | --forward-model-noise-std 0.0 \ 28 | --forward-model-batch-size 128 \ 29 | --forward-model-val-size 500 \ 30 | --forward-model-epochs 50 \ 31 | --evaluation-samples 128 \ 32 | --not-fast & done; done; wait; done 33 | -------------------------------------------------------------------------------- /scripts/coms-hopper-demo-conservatism.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | NUM_TRIALS_PER_GPU=2 3 | IFS=, read -ra DEVICES <<< "$CUDA_VISIBLE_DEVICES" 4 | 5 | for DEVICE in "${DEVICES[@]}"; do 6 | for TRIAL in $(seq $NUM_TRIALS_PER_GPU); do 7 | CUDA_VISIBLE_DEVICES=$DEVICE coms \ 8 | --logging-dir ~/coms-hopper-demo/coms-hopper-cons/COMs-HopperController-Exact-v0-$DEVICE-$TRIAL-$RANDOM \ 9 | --task HopperController-Exact-v0 \ 10 | --no-task-relabel \ 11 | --normalize-ys \ 12 | --normalize-xs \ 13 | --not-in-latent-space \ 14 | --particle-lr 0.05 \ 15 | --particle-train-gradient-steps 50 \ 16 | --particle-evaluate-gradient-steps 50 \ 17 | --particle-entropy-coefficient 0.0 \ 18 | --forward-model-activations relu \ 19 | --forward-model-activations relu \ 20 | --forward-model-hidden-size 2048 \ 21 | --no-forward-model-final-tanh \ 22 | --forward-model-lr 0.0003 \ 23 | --forward-model-alpha 0.1 \ 24 | --forward-model-alpha-lr 0.01 \ 25 | --forward-model-overestimation-limit 0.5 \ 26 | --forward-model-noise-std 0.0 \ 27 | --forward-model-batch-size 128 \ 28 | --forward-model-val-size 500 \ 29 | --forward-model-epochs 50 \ 30 | --evaluation-samples 128 \ 31 | --not-fast & done; done; wait 32 | 33 | for DEVICE in "${DEVICES[@]}"; do 34 | for TRIAL in $(seq $NUM_TRIALS_PER_GPU); do 35 | CUDA_VISIBLE_DEVICES=$DEVICE coms \ 36 | --logging-dir ~/coms-hopper-demo/coms-hopper-over/COMs-HopperController-Exact-v0-$DEVICE-$TRIAL-$RANDOM \ 37 | --task HopperController-Exact-v0 \ 38 | --no-task-relabel \ 39 | --normalize-ys \ 40 | --normalize-xs \ 41 | --not-in-latent-space \ 42 | --particle-lr 0.05 \ 43 | --particle-train-gradient-steps 50 \ 44 | --particle-evaluate-gradient-steps 50 \ 45 | --particle-entropy-coefficient 0.0 \ 46 | --forward-model-activations relu \ 47 | --forward-model-activations relu \ 48 | --forward-model-hidden-size 2048 \ 49 | --no-forward-model-final-tanh \ 50 | --forward-model-lr 0.0003 \ 51 | --forward-model-alpha 0.0 \ 52 | --forward-model-alpha-lr 0.0 \ 53 | --forward-model-overestimation-limit 10000.0 \ 54 | --forward-model-noise-std 0.0 \ 55 | --forward-model-batch-size 128 \ 56 | --forward-model-val-size 500 \ 57 | --forward-model-epochs 50 \ 58 | --evaluation-samples 128 \ 59 | --not-fast & done; done; wait 60 | -------------------------------------------------------------------------------- /scripts/coms-hopper.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | NUM_TRIALS_PER_GPU=2 3 | IFS=, read -ra DEVICES <<< "$CUDA_VISIBLE_DEVICES" 4 | for OE_LIMIT in 0.4 0.5 0.6 0.8 1.0; do 5 | for DEVICE in "${DEVICES[@]}"; do 6 | for TRIAL in $(seq $NUM_TRIALS_PER_GPU); do 7 | CUDA_VISIBLE_DEVICES=$DEVICE coms \ 8 | --logging-dir ~/coms-hopper/coms-hopper-$OE_LIMIT/COMs-HopperController-Exact-v0-$DEVICE-$TRIAL-$RANDOM \ 9 | --task HopperController-Exact-v0 \ 10 | --no-task-relabel \ 11 | --normalize-ys \ 12 | --normalize-xs \ 13 | --not-in-latent-space \ 14 | --vae-hidden-size 64 \ 15 | --vae-latent-size 256 \ 16 | --vae-activation relu \ 17 | --vae-kernel-size 3 \ 18 | --vae-num-blocks 4 \ 19 | --vae-lr 0.0003 \ 20 | --vae-beta 1.0 \ 21 | --vae-batch-size 128 \ 22 | --vae-val-size 500 \ 23 | --vae-epochs 10 \ 24 | --particle-lr 0.05 \ 25 | --particle-train-gradient-steps 50 \ 26 | --particle-evaluate-gradient-steps 50 \ 27 | --particle-entropy-coefficient 0.0 \ 28 | --forward-model-activations relu \ 29 | --forward-model-activations relu \ 30 | --forward-model-hidden-size 2048 \ 31 | --no-forward-model-final-tanh \ 32 | --forward-model-lr 0.0003 \ 33 | --forward-model-alpha 0.1 \ 34 | --forward-model-alpha-lr 0.01 \ 35 | --forward-model-overestimation-limit $OE_LIMIT \ 36 | --forward-model-noise-std 0.0 \ 37 | --forward-model-batch-size 128 \ 38 | --forward-model-val-size 500 \ 39 | --forward-model-epochs 50 \ 40 | --evaluation-samples 128 \ 41 | --fast & done; done; wait 42 | done 43 | -------------------------------------------------------------------------------- /scripts/coms-nas.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | . /packages/anaconda3/etc/profile.d/conda.shi 3 | conda activate design-baselines 4 | NUM_TRIALS_PER_GPU=2 5 | IFS=, read -ra DEVICES <<< "$CUDA_VISIBLE_DEVICES" 6 | for OE_LIMIT in 2.0; do 7 | for DEVICE in "${DEVICES[@]}"; do 8 | for TRIAL in $(seq $NUM_TRIALS_PER_GPU); do 9 | CUDA_VISIBLE_DEVICES=$DEVICE coms \ 10 | --logging-dir ./coms-nas/coms-nas-$OE_LIMIT/COMs-CIFARNAS-Exact-v0-$DEVICE-$TRIAL-$RANDOM \ 11 | --task CIFARNAS-Exact-v0 \ 12 | --no-task-relabel \ 13 | --normalize-ys \ 14 | --no-normalize-xs \ 15 | --not-in-latent-space \ 16 | --vae-hidden-size 64 \ 17 | --vae-latent-size 256 \ 18 | --vae-activation relu \ 19 | --vae-kernel-size 3 \ 20 | --vae-num-blocks 4 \ 21 | --vae-lr 0.0003 \ 22 | --vae-beta 1.0 \ 23 | --vae-batch-size 128 \ 24 | --vae-val-size 500 \ 25 | --vae-epochs 10 \ 26 | --particle-lr 2.0 \ 27 | --particle-train-gradient-steps 50 \ 28 | --particle-evaluate-gradient-steps 50 \ 29 | --particle-entropy-coefficient 0.0 \ 30 | --forward-model-activations relu \ 31 | --forward-model-activations relu \ 32 | --forward-model-hidden-size 2048 \ 33 | --no-forward-model-final-tanh \ 34 | --forward-model-lr 0.0003 \ 35 | --forward-model-alpha 0.1 \ 36 | --forward-model-alpha-lr 0.01 \ 37 | --forward-model-overestimation-limit $OE_LIMIT \ 38 | --forward-model-noise-std 0.0 \ 39 | --forward-model-batch-size 128 \ 40 | --forward-model-val-size 100 \ 41 | --forward-model-epochs 50 \ 42 | --evaluation-samples 128 \ 43 | --fast & done; done; wait 44 | done 45 | -------------------------------------------------------------------------------- /scripts/coms-superconductor.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | NUM_TRIALS_PER_GPU=2 3 | IFS=, read -ra DEVICES <<< "$CUDA_VISIBLE_DEVICES" 4 | for OE_LIMIT in 0.4 0.5 0.6 0.8 1.0; do 5 | for DEVICE in "${DEVICES[@]}"; do 6 | for TRIAL in $(seq $NUM_TRIALS_PER_GPU); do 7 | CUDA_VISIBLE_DEVICES=$DEVICE coms \ 8 | --logging-dir ~/coms-superconductor/coms-superconductor-$OE_LIMIT/COMs-Superconductor-RandomForest-v0-$DEVICE-$TRIAL-$RANDOM \ 9 | --task Superconductor-RandomForest-v0 \ 10 | --no-task-relabel \ 11 | --normalize-ys \ 12 | --normalize-xs \ 13 | --not-in-latent-space \ 14 | --vae-hidden-size 64 \ 15 | --vae-latent-size 256 \ 16 | --vae-activation relu \ 17 | --vae-kernel-size 3 \ 18 | --vae-num-blocks 4 \ 19 | --vae-lr 0.0003 \ 20 | --vae-beta 1.0 \ 21 | --vae-batch-size 128 \ 22 | --vae-val-size 500 \ 23 | --vae-epochs 10 \ 24 | --particle-lr 0.05 \ 25 | --particle-train-gradient-steps 50 \ 26 | --particle-evaluate-gradient-steps 50 \ 27 | --particle-entropy-coefficient 0.0 \ 28 | --forward-model-activations relu \ 29 | --forward-model-activations relu \ 30 | --forward-model-hidden-size 2048 \ 31 | --no-forward-model-final-tanh \ 32 | --forward-model-lr 0.0003 \ 33 | --forward-model-alpha 0.1 \ 34 | --forward-model-alpha-lr 0.01 \ 35 | --forward-model-overestimation-limit $OE_LIMIT \ 36 | --forward-model-noise-std 0.0 \ 37 | --forward-model-batch-size 128 \ 38 | --forward-model-val-size 500 \ 39 | --forward-model-epochs 50 \ 40 | --evaluation-samples 128 \ 41 | --fast & done; done; wait 42 | done 43 | -------------------------------------------------------------------------------- /scripts/coms-tf-bind-10.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | . /packages/anaconda3/etc/profile.d/conda.shi 3 | conda activate design-baselines 4 | NUM_TRIALS_PER_GPU=2 5 | IFS=, read -ra DEVICES <<< "$CUDA_VISIBLE_DEVICES" 6 | for OE_LIMIT in 2.0; do 7 | for DEVICE in "${DEVICES[@]}"; do 8 | for TRIAL in $(seq $NUM_TRIALS_PER_GPU); do 9 | CUDA_VISIBLE_DEVICES=$DEVICE coms \ 10 | --logging-dir ./coms-tf-bind-10/coms-tf-bind-10-$OE_LIMIT/COMs-TFBind10-Exact-v0-$DEVICE-$TRIAL-$RANDOM \ 11 | --task TFBind10-Exact-v0 \ 12 | --no-task-relabel \ 13 | --task-max-samples 10000 \ 14 | --task-distribution uniform \ 15 | --normalize-ys \ 16 | --no-normalize-xs \ 17 | --not-in-latent-space \ 18 | --vae-hidden-size 64 \ 19 | --vae-latent-size 256 \ 20 | --vae-activation relu \ 21 | --vae-kernel-size 3 \ 22 | --vae-num-blocks 4 \ 23 | --vae-lr 0.0003 \ 24 | --vae-beta 1.0 \ 25 | --vae-batch-size 128 \ 26 | --vae-val-size 500 \ 27 | --vae-epochs 10 \ 28 | --particle-lr 2.0 \ 29 | --particle-train-gradient-steps 50 \ 30 | --particle-evaluate-gradient-steps 50 \ 31 | --particle-entropy-coefficient 0.0 \ 32 | --forward-model-activations relu \ 33 | --forward-model-activations relu \ 34 | --forward-model-hidden-size 2048 \ 35 | --no-forward-model-final-tanh \ 36 | --forward-model-lr 0.0003 \ 37 | --forward-model-alpha 0.1 \ 38 | --forward-model-alpha-lr 0.01 \ 39 | --forward-model-overestimation-limit $OE_LIMIT \ 40 | --forward-model-noise-std 0.0 \ 41 | --forward-model-batch-size 128 \ 42 | --forward-model-val-size 500 \ 43 | --forward-model-epochs 50 \ 44 | --evaluation-samples 128 \ 45 | --fast & done; done; wait 46 | done 47 | -------------------------------------------------------------------------------- /scripts/coms-tf-bind-8.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | NUM_TRIALS_PER_GPU=2 3 | IFS=, read -ra DEVICES <<< "$CUDA_VISIBLE_DEVICES" 4 | for OE_LIMIT in 2.0; do 5 | for DEVICE in "${DEVICES[@]}"; do 6 | for TRIAL in $(seq $NUM_TRIALS_PER_GPU); do 7 | CUDA_VISIBLE_DEVICES=$DEVICE coms \ 8 | --logging-dir ~/coms-tf-bind-8/coms-tf-bind-8-$OE_LIMIT/COMs-TFBind8-Exact-v0-$DEVICE-$TRIAL-$RANDOM \ 9 | --task TFBind8-Exact-v0 \ 10 | --no-task-relabel \ 11 | --normalize-ys \ 12 | --no-normalize-xs \ 13 | --not-in-latent-space \ 14 | --vae-hidden-size 64 \ 15 | --vae-latent-size 256 \ 16 | --vae-activation relu \ 17 | --vae-kernel-size 3 \ 18 | --vae-num-blocks 4 \ 19 | --vae-lr 0.0003 \ 20 | --vae-beta 1.0 \ 21 | --vae-batch-size 128 \ 22 | --vae-val-size 500 \ 23 | --vae-epochs 10 \ 24 | --particle-lr 2.0 \ 25 | --particle-train-gradient-steps 50 \ 26 | --particle-evaluate-gradient-steps 50 \ 27 | --particle-entropy-coefficient 0.0 \ 28 | --forward-model-activations relu \ 29 | --forward-model-activations relu \ 30 | --forward-model-hidden-size 2048 \ 31 | --no-forward-model-final-tanh \ 32 | --forward-model-lr 0.0003 \ 33 | --forward-model-alpha 0.1 \ 34 | --forward-model-alpha-lr 0.01 \ 35 | --forward-model-overestimation-limit $OE_LIMIT \ 36 | --forward-model-noise-std 0.0 \ 37 | --forward-model-batch-size 128 \ 38 | --forward-model-val-size 500 \ 39 | --forward-model-epochs 50 \ 40 | --evaluation-samples 128 \ 41 | --fast & done; done; wait 42 | done 43 | -------------------------------------------------------------------------------- /scripts/coms-utr-ablate-conservatism.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | NUM_TRIALS_PER_GPU=2 3 | IFS=, read -ra DEVICES <<< "$CUDA_VISIBLE_DEVICES" 4 | 5 | for OE_LIMIT in 0.5 1.0 2.0 4.0 8.0; do 6 | for DEVICE in "${DEVICES[@]}"; do 7 | for TRIAL in $(seq $NUM_TRIALS_PER_GPU); do 8 | CUDA_VISIBLE_DEVICES=$DEVICE coms \ 9 | --logging-dir ~/coms-utr-ablate/coms-utr-$OE_LIMIT/COMs-UTR-ResNet-v0-$DEVICE-$TRIAL-$RANDOM \ 10 | --task UTR-ResNet-v0 \ 11 | --task-relabel \ 12 | --normalize-ys \ 13 | --normalize-xs \ 14 | --not-in-latent-space \ 15 | --particle-lr 2.0 \ 16 | --particle-train-gradient-steps 50 \ 17 | --particle-evaluate-gradient-steps 50 \ 18 | --particle-entropy-coefficient 0.0 \ 19 | --forward-model-activations relu \ 20 | --forward-model-activations relu \ 21 | --forward-model-hidden-size 2048 \ 22 | --no-forward-model-final-tanh \ 23 | --forward-model-lr 0.0003 \ 24 | --forward-model-alpha 0.1 \ 25 | --forward-model-alpha-lr 0.01 \ 26 | --forward-model-overestimation-limit $OE_LIMIT \ 27 | --forward-model-noise-std 0.0 \ 28 | --forward-model-batch-size 128 \ 29 | --forward-model-val-size 500 \ 30 | --forward-model-epochs 50 \ 31 | --evaluation-samples 128 \ 32 | --not-fast & done; done; wait 33 | -------------------------------------------------------------------------------- /scripts/coms-utr-demo-conservatism.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | NUM_TRIALS_PER_GPU=2 3 | IFS=, read -ra DEVICES <<< "$CUDA_VISIBLE_DEVICES" 4 | 5 | for DEVICE in "${DEVICES[@]}"; do 6 | for TRIAL in $(seq $NUM_TRIALS_PER_GPU); do 7 | CUDA_VISIBLE_DEVICES=$DEVICE coms \ 8 | --logging-dir ~/coms-utr-demo/coms-utr-cons/COMs-UTR-ResNet-v0-$DEVICE-$TRIAL-$RANDOM \ 9 | --task UTR-ResNet-v0 \ 10 | --task-relabel \ 11 | --normalize-ys \ 12 | --normalize-xs \ 13 | --not-in-latent-space \ 14 | --particle-lr 2.0 \ 15 | --particle-train-gradient-steps 50 \ 16 | --particle-evaluate-gradient-steps 50 \ 17 | --particle-entropy-coefficient 0.0 \ 18 | --forward-model-activations relu \ 19 | --forward-model-activations relu \ 20 | --forward-model-hidden-size 2048 \ 21 | --no-forward-model-final-tanh \ 22 | --forward-model-lr 0.0003 \ 23 | --forward-model-alpha 0.1 \ 24 | --forward-model-alpha-lr 0.01 \ 25 | --forward-model-overestimation-limit 2.0 \ 26 | --forward-model-noise-std 0.0 \ 27 | --forward-model-batch-size 128 \ 28 | --forward-model-val-size 500 \ 29 | --forward-model-epochs 50 \ 30 | --evaluation-samples 128 \ 31 | --not-fast & done; done; wait 32 | 33 | for DEVICE in "${DEVICES[@]}"; do 34 | for TRIAL in $(seq $NUM_TRIALS_PER_GPU); do 35 | CUDA_VISIBLE_DEVICES=$DEVICE coms \ 36 | --logging-dir ~/coms-utr-demo/coms-utr-over/COMs-UTR-ResNet-v0-$DEVICE-$TRIAL-$RANDOM \ 37 | --task UTR-ResNet-v0 \ 38 | --task-relabel \ 39 | --normalize-ys \ 40 | --normalize-xs \ 41 | --not-in-latent-space \ 42 | --particle-lr 2.0 \ 43 | --particle-train-gradient-steps 50 \ 44 | --particle-evaluate-gradient-steps 50 \ 45 | --particle-entropy-coefficient 0.0 \ 46 | --forward-model-activations relu \ 47 | --forward-model-activations relu \ 48 | --forward-model-hidden-size 2048 \ 49 | --no-forward-model-final-tanh \ 50 | --forward-model-lr 0.0003 \ 51 | --forward-model-alpha 0.0 \ 52 | --forward-model-alpha-lr 0.0 \ 53 | --forward-model-overestimation-limit 10000.0 \ 54 | --forward-model-noise-std 0.0 \ 55 | --forward-model-batch-size 128 \ 56 | --forward-model-val-size 500 \ 57 | --forward-model-epochs 50 \ 58 | --evaluation-samples 128 \ 59 | --not-fast & done; done; wait 60 | -------------------------------------------------------------------------------- /scripts/coms-utr.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | NUM_TRIALS_PER_GPU=2 3 | IFS=, read -ra DEVICES <<< "$CUDA_VISIBLE_DEVICES" 4 | for OE_LIMIT in 2.0; do 5 | for DEVICE in "${DEVICES[@]}"; do 6 | for TRIAL in $(seq $NUM_TRIALS_PER_GPU); do 7 | CUDA_VISIBLE_DEVICES=$DEVICE coms \ 8 | --logging-dir ~/coms-utr/coms-utr-$OE_LIMIT/COMs-UTR-Transformer-v0-$DEVICE-$TRIAL-$RANDOM \ 9 | --task UTR-ResNet-v0 \ 10 | --task-relabel \ 11 | --normalize-ys \ 12 | --no-normalize-xs \ 13 | --not-in-latent-space \ 14 | --vae-hidden-size 64 \ 15 | --vae-latent-size 256 \ 16 | --vae-activation relu \ 17 | --vae-kernel-size 3 \ 18 | --vae-num-blocks 4 \ 19 | --vae-lr 0.0003 \ 20 | --vae-beta 1.0 \ 21 | --vae-batch-size 128 \ 22 | --vae-val-size 500 \ 23 | --vae-epochs 10 \ 24 | --particle-lr 2.0 \ 25 | --particle-train-gradient-steps 50 \ 26 | --particle-evaluate-gradient-steps 50 \ 27 | --particle-entropy-coefficient 0.0 \ 28 | --forward-model-activations relu \ 29 | --forward-model-activations relu \ 30 | --forward-model-hidden-size 2048 \ 31 | --no-forward-model-final-tanh \ 32 | --forward-model-lr 0.0003 \ 33 | --forward-model-alpha 0.1 \ 34 | --forward-model-alpha-lr 0.01 \ 35 | --forward-model-overestimation-limit $OE_LIMIT \ 36 | --forward-model-noise-std 0.0 \ 37 | --forward-model-batch-size 128 \ 38 | --forward-model-val-size 500 \ 39 | --forward-model-epochs 50 \ 40 | --evaluation-samples 128 \ 41 | --fast & done; done; wait 42 | done 43 | -------------------------------------------------------------------------------- /scripts/launch_chembl.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | CUDA_VISIBLE_DEVICES=0,1,2,3 bash $HOME/design-baselines/scripts/coms-chembl.sh 3 | -------------------------------------------------------------------------------- /scripts/launch_demo.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | CUDA_VISIBLE_DEVICES=0,1,2,3 bash $HOME/design-baselines/scripts/coms-ant.sh 3 | CUDA_VISIBLE_DEVICES=0,1,2,3 bash $HOME/design-baselines/scripts/coms-dkitty.sh 4 | CUDA_VISIBLE_DEVICES=0,1,2,3 bash $HOME/design-baselines/scripts/coms-hopper.sh 5 | CUDA_VISIBLE_DEVICES=0,1,2,3 bash $HOME/design-baselines/scripts/coms-superconductor.sh 6 | CUDA_VISIBLE_DEVICES=0,1,2,3 bash $HOME/design-baselines/scripts/coms-gfp.sh 7 | CUDA_VISIBLE_DEVICES=0,1,2,3 bash $HOME/design-baselines/scripts/coms-utr.sh 8 | CUDA_VISIBLE_DEVICES=0,1,2,3 bash $HOME/design-baselines/scripts/coms-tf-bind-8.sh 9 | CUDA_VISIBLE_DEVICES=0,1,2,3 bash $HOME/design-baselines/scripts/coms-chembl.sh 10 | -------------------------------------------------------------------------------- /scripts/launch_nas.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | CUDA_VISIBLE_DEVICES=0,1,2,3 bash $HOME/design-baselines/scripts/coms-nas.sh 3 | -------------------------------------------------------------------------------- /scripts/launch_tf_bind_10.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | CUDA_VISIBLE_DEVICES=0,1,2,3 bash $HOME/design-baselines/scripts/coms-tf-bind-10.sh 3 | -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | from setuptools import find_packages 2 | from setuptools import setup 3 | 4 | 5 | setup( 6 | name='design-baselines', 7 | description='Baselines for Model-Based Optimization', 8 | license='MIT', 9 | version='0.1', 10 | zip_safe=True, 11 | include_package_data=True, 12 | packages=find_packages(), 13 | entry_points={'console_scripts': ( 14 | 'design-baselines=design_baselines.cli:cli', 15 | 16 | 'coms=design_baselines.coms_cleaned:coms_cleaned', 17 | 18 | 'coms-cleaned=design_baselines.' 19 | 'coms_cleaned.experiments:cli', 20 | 'coms-original=design_baselines.' 21 | 'coms_original.experiments:cli', 22 | 23 | 'gradient-ascent=design_baselines.' 24 | 'gradient_ascent.experiments:cli', 25 | 'gradient-ascent-relabelled=design_baselines.' 26 | 'gradient_ascent.relabel_experiments:cli', 27 | 'gradient-ascent-ablate-distribution=design_baselines.' 28 | 'gradient_ascent.distribution_experiments:cli', 29 | 30 | 'gradient-ascent-min-ensemble=design_baselines.' 31 | 'gradient_ascent.min_ensemble_experiments:cli', 32 | 'gradient-ascent-min-ensemble-relabelled=design_baselines.' 33 | 'gradient_ascent.relabel_min_ensemble_experiments:cli', 34 | 'gradient-ascent-min-ensemble-ablate-distribution=design_baselines.' 35 | 'gradient_ascent.distribution_min_ensemble_experiments:cli', 36 | 37 | 'gradient-ascent-mean-ensemble=design_baselines.' 38 | 'gradient_ascent.mean_ensemble_experiments:cli', 39 | 'gradient-ascent-mean-ensemble-relabelled=design_baselines.' 40 | 'gradient_ascent.relabel_mean_ensemble_experiments:cli', 41 | 'gradient-ascent-mean-ensemble-ablate-distribution=design_baselines.' 42 | 'gradient_ascent.distribution_mean_ensemble_experiments:cli', 43 | 44 | 'mins=design_baselines.' 45 | 'mins.experiments:cli', 46 | 'mins-relabelled=design_baselines.' 47 | 'mins.relabel_experiments:cli', 48 | 'mins-ablate-distribution=design_baselines.' 49 | 'mins.distribution_experiments:cli', 50 | 51 | 'cbas=design_baselines.' 52 | 'cbas.experiments:cli', 53 | 'cbas-relabelled=design_baselines.' 54 | 'cbas.relabel_experiments:cli', 55 | 'cbas-ablate-distribution=design_baselines.' 56 | 'cbas.distribution_experiments:cli', 57 | 58 | 'autofocused-cbas=design_baselines.' 59 | 'autofocused_cbas.experiments:cli', 60 | 'autofocused-cbas-relabelled=design_baselines.' 61 | 'autofocused_cbas.relabel_experiments:cli', 62 | 'autofocused-cbas-ablate-distribution=design_baselines.' 63 | 'autofocused_cbas.distribution_experiments:cli', 64 | 65 | 'cma-es=design_baselines.' 66 | 'cma_es.experiments:cli', 67 | 'cma-es-relabelled=design_baselines.' 68 | 'cma_es.relabel_experiments:cli', 69 | 'cma-es-ablate-distribution=design_baselines.' 70 | 'cma_es.distribution_experiments:cli', 71 | 72 | 'bo-qei=design_baselines.' 73 | 'bo_qei.experiments:cli', 74 | 'bo-qei-relabelled=design_baselines.' 75 | 'bo_qei.relabel_experiments:cli', 76 | 'bo-qei-ablate-distribution=design_baselines.' 77 | 'bo_qei.distribution_experiments:cli', 78 | 79 | 'reinforce=design_baselines.' 80 | 'reinforce.experiments:cli', 81 | 'reinforce-relabelled=design_baselines.' 82 | 'reinforce.relabel_experiments:cli', 83 | 'reinforce-ablate-distribution=design_baselines.' 84 | 'reinforce.distribution_experiments:cli', 85 | 86 | 'online-reinforce=design_baselines.' 87 | 'reinforce.online_experiments:cli', 88 | )}) 89 | --------------------------------------------------------------------------------