Project: 011_Standard_Gaussian_Process_is_All_You_Need_for_High | Function: BO_loop_GP (iclr)

system: You are an expert Python programmer implementing algorithms from scientific papers.
Complete the masked function based on the algorithm description.
Follow Python best practices and maintain consistent code style.
Return **only** the complete Python function code.

user: Complete this Python function:

**DESCRIPTION:** Main Bayesian optimization loop implementing the paper's approach with UCB acquisition function and robust GP training for high-dimensional optimization

**KEY TERMS:** {
  "UpperConfidenceBound": "UCB acquisition function",
  "beta=1.5": "λ = 1.5 (confidence parameter)",
  "if_matern": "Matérn vs RBF kernel selection",
  "set_ls": "robust initialization flag",
  "optimize_acqf": "acquisition function optimization"
}

**CONTEXT CODE:** from botorch.acquisition.analytic import ExpectedImprovement, UpperConfidenceBound, LogExpectedImprovement
from botorch.optim import optimize_acqf
from baselines.GP import GP_Wrapper, GP_MAP_Wrapper, Vanilla_GP_Wrapper
from data import *
import time
from infras.randutils import *



def BO_loop_GP(func_name, dataset, seed, num_step=200, beta=1.5, if_ard=False, if_softplus=True, acqf_type="UCB", if_matern=True, set_ls=False, device="cpu"):

    raise NotImplementedError("Function implementation has been masked for testing")

def Vanilla_BO_loop(func_name, dataset, seed, num_step=200):
    """
    Our implementation for Vanilla BO
    """
    best_y = []
    time_list = []
    dim = dataset.func.dims

    for i in range(1, num_step + 1):
        start_time = time.time()
        X, Y = dataset.get_data(normalize=True)
        best_y_before = dataset.get_curr_max_unnormed()
        model = Vanilla_GP_Wrapper(X, Y)
        model.train_model()

        ls = model.gp_model.covar_module.base_kernel.lengthscale
        print(f"ls mean: {ls.mean()}, ls std: {ls.std()}, max: {ls.max()}, min: {ls.min()}")

        acqf = LogExpectedImprovement(model=model.gp_model, best_f=Y.max())
        new_x, _ = optimize_acqf(
            acq_function=acqf,
            bounds=torch.tensor([[0.0] * dim, [1.0] * dim]),
            q=1,
            num_restarts=10,
            raw_samples=1000,
            options={},
        )

        end_time = time.time()
        time_used = end_time - start_time
        time_list.append(time_used)
        dataset.add(new_x)
        best_y_after = dataset.get_curr_max_unnormed()

        print(
            f"Seed: {seed} --- At itr: {i}: best value before={best_y_before}, best value after={best_y_after}, current query: {dataset.y[-1]}",
            flush=True)
        best_y.append(best_y_before)
    return best_y, time_list


def BO_loop_GP_MAP(func_name, dataset, seed, num_step=200, beta=1.5, if_ard=True, optim_type="LBFGS", acqf_type="UCB",
                   ls_prior_type="Gamma", set_ls=False, if_matern=False, device="cpu"):
    best_y = []
    time_list = []
    dim = dataset.func.dims
    for i in range(1, num_step + 1):
        start_time = time.time()
        X, Y = dataset.get_data(normalize=True)
        X = X.to(device)
        Y = Y.to(device)
        best_y_before = dataset.get_curr_max_unnormed()
        model = GP_MAP_Wrapper(X, Y, if_ard=if_ard, if_matern=if_matern, optim_type=optim_type,
                               ls_prior_type=ls_prior_type, device=device, set_ls=set_ls)
        model.train_model()

        if acqf_type == "UCB":
            acqf = UpperConfidenceBound(model=model.gp_model, beta=beta, maximize=True).to(device)
        elif acqf_type == "EI":
            acqf = ExpectedImprovement(model=model.gp_model, best_f=Y.max()).to(device)
        else:
            raise NotImplementedError

        new_x, _ = optimize_acqf(
                acq_function=acqf,
                bounds=torch.tensor([[0.0] * dim, [1.0] * dim]),
                q=1,
                num_restarts=10,
                raw_samples=1000,
                options={},
            )

        end_time = time.time()
        time_used = end_time - start_time

        time_list.append(time_used)
        dataset.add(new_x)
        best_y_after = dataset.get_curr_max_unnormed()
        print(
            f"Seed: {seed} --- At itr: {i}: best value before={best_y_before}, best value after={best_y_after}, current query: {dataset.y[-1]}",
            flush=True)
        best_y.append(best_y_before)
    return best_y, time_list




Complete the function implementation replacing "raise NotImplementedError".