optim_PhC_so.py

"""Bayesian optimization for optimizing PCSELs (using OpenAI Gym).
#Renjie Li, May 2023, NOEL CUHKSZ.
"""
#2023.5.15: this version increases N_batch and beta in UCB.

#import sys
import gym
from gym.envs.registration import register
#import math
#import random
import numpy as np
from botorch.models import SingleTaskGP
from botorch.fit import fit_gpytorch_mll
from botorch.utils import standardize
from gpytorch.mlls import ExactMarginalLogLikelihood
from botorch.acquisition import UpperConfidenceBound
from botorch.optim import optimize_acqf
from botorch.models.transforms import Normalize, Standardize
#import matplotlib
#import matplotlib.pyplot as plt
import torch
#import torch.nn as nn
#import torch.optim as optim 
#import torch.nn.functional as F
#from torch.utils.tensorboard import SummaryWriter
from tensorboardX import SummaryWriter
#import torchvision.transforms as T
import logging
from itertools import count
from datetime import datetime, timezone
import time

torch.set_printoptions(precision=10)

logger = logging.getLogger(__name__)

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# register the env with gym
register(
    id='Fdtd_NB-v0',
    entry_point='envs:FdtdEnv',
    max_episode_steps=150,
    reward_threshold=250.0,
)

writer = SummaryWriter()  # log the training process

# instantiate the fdtd env
env = gym.make('Fdtd_NB-v0').unwrapped

state = env.reset()


def generate_initial_data():
    # generate training data
    #input variable: state
    #train_X = torch.zeros(1, 8, device=device, dtype=torch.float64)  #fixed zero initialization
    train_X = torch.tensor([[-818.275115, -2.524478, -4.409965, -88.682263, 0.04145147610,
                             0.0934398093, 0.231437257, 27.566375]], device=device, dtype=torch.float64)  #fixed optimal initialization
    # x1 = -200*torch.rand(1, 1, device=device, dtype=torch.float64) + 100   #random initialization
    # x2 = -0.3*torch.rand(1, 1, device=device, dtype=torch.float64) + 0.15
    # x3 = -0.6*torch.rand(1, 1, device=device, dtype=torch.float64) + 0.3
    # x4 = -2000*torch.rand(1, 1, device=device, dtype=torch.float64) + 1000
    # train_X = torch.hstack((x4, x1, x1, x1, x2, x2, x3, x1))

    score = env.step(train_X.tolist()[0])

    # output variable: score
    train_Y = torch.tensor([score], device=device)

    train_Y = torch.reshape(train_Y, (1,1)) #make dim (1,1)

    # train_Y = standardize(train_Y)

    return train_X, train_Y

def initialize_model(train_X, train_Y, state_dict=None):
    # define models for objective and constraint
    gp = SingleTaskGP(train_X, train_Y, outcome_transform=Standardize(m=train_Y.shape[-1]), input_transform=Normalize(d=train_X.shape[-1]))
    mll = ExactMarginalLogLikelihood(gp.likelihood, gp)

    # load state dict if it is passed
    if state_dict is not None:
        gp.load_state_dict(state_dict)

    return mll, gp


def optimize_acqf_and_get_observation(acq_func):
    """Optimizes the acquisition function, and returns a new candidate and a noisy observation."""
    global score
    high = [1000.0, 100., 100., 100., 0.15, 0.15, 0.3, 100.]
    low = [-1000.0, -100., -100., -100., -0.15, -0.15, -0.3, -100.]
    bounds = torch.tensor([low, high], dtype=torch.float64)
    candidate, acq_value = optimize_acqf(
        acq_function = acq_func, bounds=bounds, q=1, num_restarts=5, raw_samples=20,
    )

    # observe new values
    new_x = candidate.detach()
    #new_x = torch.round(new_x)
    score = env.step(new_x.tolist()[0])
    new_y = torch.tensor([score], device=device)
    new_y = torch.reshape(new_y, (1,1)) #make dim (1,1)

    #print('score: {:.5f}, state: {}\n'.format(score, new_x.tolist()[0]))
    #new_y = standardize(new_y)
    return new_x, new_y, acq_value


N_TRIALS = 60
N_BATCH = 400

verbose = True

#save overall best solutions
best_observed_all_ei = []

#save data
input = torch.tensor([])
objective = torch.tensor([])

steps_done = 0

# average over multiple trials
for trial in range(1, N_TRIALS + 1):
    utc_dt = datetime.now(timezone.utc)
    print("\nLocal time {}".format(utc_dt.astimezone().isoformat()))

    print(f"\nTrial {trial:>2} of {N_TRIALS} \n", end="")
    best_observed_ei = []

    env.reset()

    # call helper functions to generate initial training data and initialize model
    (
        train_x_ei,
        train_obj_ei,
    ) = generate_initial_data()

    print('\nx_initial: {}, f_initial: {:.5f}\n'.format(train_x_ei.tolist()[0], train_obj_ei.item()))

    mll_ei, model_ei = initialize_model(train_x_ei, train_obj_ei)

    best_observed_ei.append(train_obj_ei)

    # run N_BATCH rounds of BayesOpt after the initial random batch
    for iteration in range(1, N_BATCH + 1):
        print('\nStarting step No.{}'.format(iteration))

        steps_done += 1

        t0 = time.monotonic()

        # fit the models
        fit_gpytorch_mll(mll_ei)

        #acquisition function
        qEI = UpperConfidenceBound(model_ei, beta=2.0) #increase beta for more exploration

        # optimize and get new observation
        new_x_ei, new_obj_ei, acq = optimize_acqf_and_get_observation(qEI)

        print('\nx_candidate: {}, f_score: {:.5f}, acquisition: {:.5f}\n'.format(new_x_ei.tolist()[0], new_obj_ei.item(), acq.item()))

        writer.add_scalar('BO/scores', new_obj_ei.item(), steps_done)
        writer.add_scalar('BO/acq value', acq.item(), steps_done)

        # update training points
        train_x_ei = torch.cat([train_x_ei, new_x_ei])
        train_obj_ei = torch.cat([train_obj_ei, new_obj_ei])

        # update progress
        best_value_ei = train_obj_ei.max().item()
        best_observed_ei.append(best_value_ei)

        # reinitialize the models so they are ready for fitting on next iteration
        # use the current state dict to speed up fitting
        mll_ei, model_ei = initialize_model(
            train_x_ei,
            train_obj_ei,
            model_ei.state_dict(),
        )

        t1 = time.monotonic()

        if iteration % 50 == 0:
            input = torch.cat([input, train_x_ei])
            objective = torch.cat([objective, train_obj_ei])
            torch.save(input, 'input.pt')
            torch.save(objective, 'objective.pt')

    if verbose:
        print(
                f"\nBatch {iteration:>2}: best_score = "
                f"({best_value_ei:>4.5f}),"
                f"time = {t1-t0:>4.2f}.",
                end="",
            )
    else:
        print(".", end="")

    writer.add_scalar('BO/best score', best_value_ei, trial)

    best_observed_all_ei.append(best_observed_ei)

    print(best_observed_ei)

writer.close()