def testContours(self):
exp = get_branin_experiment(with_str_choice_param=True,
with_batch=True)
exp.trials[0].run()
model = Models.BOTORCH(
# Model bridge kwargs
experiment=exp,
data=exp.fetch_data(),
)
# Assert that each type of plot can be constructed successfully
plot = plot_contour_plotly(model, model.parameters[0],
model.parameters[1],
list(model.metric_names)[0])
self.assertIsInstance(plot, go.Figure)
plot = interact_contour_plotly(model, list(model.metric_names)[0])
self.assertIsInstance(plot, go.Figure)
plot = interact_contour(model, list(model.metric_names)[0])
self.assertIsInstance(plot, AxPlotConfig)
plot = plot = plot_contour(model, model.parameters[0],
model.parameters[1],
list(model.metric_names)[0])
self.assertIsInstance(plot, AxPlotConfig)
# Make sure all parameters and metrics are displayed in tooltips
tooltips = list(exp.parameters.keys()) + list(exp.metrics.keys())
for d in plot.data["data"]:
# Only check scatter plots hoverovers
if d["type"] != "scatter":
continue
for text in d["text"]:
for tt in tooltips:
self.assertTrue(tt in text)
def get_modelbridge() -> ModelBridge:
exp = get_branin_experiment(with_batch=True)
exp.trials[0].run()
return Models.BOTORCH(
# Model bridge kwargs
experiment=exp,
data=exp.fetch_data(),
)
def test_get_standard_plots(self):
exp = get_branin_experiment()
self.assertEqual(
len(
get_standard_plots(experiment=exp,
model=get_generation_strategy().model)),
0,
)
exp = get_branin_experiment(with_batch=True, minimize=True)
exp.trials[0].run()
gs = choose_generation_strategy(search_space=exp.search_space)
gs._model = Models.BOTORCH(experiment=exp, data=exp.fetch_data())
plots = get_standard_plots(experiment=exp, model=gs.model)
self.assertEqual(len(plots), 5)
self.assertTrue(all(isinstance(plot, go.Figure) for plot in plots))
exp = get_branin_experiment_with_multi_objective(with_batch=True)
exp.trials[0].run()
gs = choose_generation_strategy(
search_space=exp.search_space,
optimization_config=exp.optimization_config)
gs._model = Models.BOTORCH(experiment=exp, data=exp.fetch_data())
plots = get_standard_plots(experiment=exp, model=gs.model)
self.assertEqual(len(plots), 6)
def test_cross_validation(self):
exp = get_branin_experiment(with_batch=True)
exp.trials[0].run()
model = Models.BOTORCH(
# Model bridge kwargs
experiment=exp,
data=exp.fetch_data(),
)
cv = cross_validate(model)
# Assert that each type of plot can be constructed successfully
plot = interact_cross_validation_plotly(cv)
self.assertIsInstance(plot, go.Figure)
plot = interact_cross_validation(cv)
self.assertIsInstance(plot, AxPlotConfig)
def testSlices(self):
exp = get_branin_experiment(with_batch=True)
exp.trials[0].run()
model = Models.BOTORCH(
# Model bridge kwargs
experiment=exp,
data=exp.fetch_data(),
)
# Assert that each type of plot can be constructed successfully
plot = plot_slice_plotly(model, model.parameters[0],
list(model.metric_names)[0])
self.assertIsInstance(plot, go.Figure)
plot = interact_slice_plotly(model)
self.assertIsInstance(plot, go.Figure)
plot = plot_slice(model, model.parameters[0],
list(model.metric_names)[0])
self.assertIsInstance(plot, AxPlotConfig)
plot = interact_slice(model)
self.assertIsInstance(plot, AxPlotConfig)
def get_GPEI(
experiment: Experiment,
data: Data,
search_space: Optional[SearchSpace] = None,
dtype: torch.dtype = torch.double,
device: torch.device = DEFAULT_TORCH_DEVICE,
) -> TorchModelBridge:
"""Instantiates a GP model that generates points with EI."""
if data.df.empty: # pragma: no cover
raise ValueError("GP+EI BotorchModel requires non-empty data.")
return checked_cast(
TorchModelBridge,
Models.BOTORCH(
experiment=experiment,
data=data,
search_space=search_space or experiment.search_space,
torch_dtype=dtype,
torch_device=device,
),
)
def testTraces(self):
exp = get_branin_experiment(with_batch=True)
exp.trials[0].run()
model = Models.BOTORCH(
# Model bridge kwargs
experiment=exp,
data=exp.fetch_data(),
)
# Assert that each type of plot can be constructed successfully
plot = optimization_trace_single_method_plotly(
np.array([[1, 2, 3], [4, 5, 6]]),
list(model.metric_names)[0],
optimization_direction="minimize",
)
self.assertIsInstance(plot, go.Figure)
plot = optimization_trace_single_method(
np.array([[1, 2, 3], [4, 5, 6]]),
list(model.metric_names)[0],
optimization_direction="minimize",
)
self.assertIsInstance(plot, AxPlotConfig)
def get_botorch(
experiment: Experiment,
data: Data,
search_space: Optional[SearchSpace] = None,
dtype: torch.dtype = torch.double,
device: torch.device = DEFAULT_TORCH_DEVICE,
transforms: List[Type[Transform]] = Cont_X_trans + Y_trans,
model_constructor: TModelConstructor = get_and_fit_model, # pyre-ignore[9]
model_predictor: TModelPredictor = predict_from_model,
acqf_constructor: TAcqfConstructor = get_NEI, # pyre-ignore[9]
acqf_optimizer: TOptimizer = scipy_optimizer, # pyre-ignore[9]
refit_on_cv: bool = False,
refit_on_update: bool = True,
optimization_config: Optional[OptimizationConfig] = None,
) -> TorchModelBridge:
"""Instantiates a BotorchModel."""
if data.df.empty: # pragma: no cover
raise ValueError("`BotorchModel` requires non-empty data.")
logger.info(
"Factory functions (like `get_botorch`) will soon be deprecated). Use "
"the model registry instead (`Models.BOTORCH(...)`).")
return checked_cast(
TorchModelBridge,
Models.BOTORCH(
experiment=experiment,
data=data,
search_space=search_space or experiment.search_space,
torch_dtype=dtype,
torch_device=device,
transforms=transforms,
model_constructor=model_constructor,
model_predictor=model_predictor,
acqf_constructor=acqf_constructor,
acqf_optimizer=acqf_optimizer,
refit_on_cv=refit_on_cv,
refit_on_update=refit_on_update,
optimization_config=optimization_config,
),
)
def get_GPEI(
experiment: Experiment,
data: Data,
search_space: Optional[SearchSpace] = None,
dtype: torch.dtype = torch.double,
device: torch.device = DEFAULT_TORCH_DEVICE,
) -> TorchModelBridge:
"""Instantiates a GP model that generates points with EI."""
if data.df.empty: # pragma: no cover
raise ValueError("GP+EI BotorchModel requires non-empty data.")
logger.info(
"Factory functions (like `get_GPEI`) will soon be deprecated). Use "
"the model registry instead (`Models.GPEI(...)`).")
return checked_cast(
TorchModelBridge,
Models.BOTORCH(
experiment=experiment,
data=data,
search_space=search_space or experiment.search_space,
torch_dtype=dtype,
torch_device=device,
),
)
def get_MOO_RS(
experiment: Experiment,
data: Data,
ref_point: Optional[List[float]] = None,
search_space: Optional[SearchSpace] = None,
dtype: torch.dtype = torch.double,
device: torch.device = DEFAULT_TORCH_DEVICE,
) -> TorchModelBridge:
"""Instantiates a Linear Random Scalarization multi-objective model.
Chooses a different random linear scalarization of the objectives
for generating each new candidate arm. This will only explore the
convex hull of the pareto frontier.
"""
# pyre-ignore: [16] `Optional` has no attribute `objective`.
if not isinstance(experiment.optimization_config.objective,
MultiObjective):
raise ValueError(
"Multi-Objective optimization requires multiple objectives")
if data.df.empty:
raise ValueError("MultiObjectiveOptimization requires non-empty data.")
return checked_cast(
TorchModelBridge,
Models.BOTORCH(
experiment=experiment,
data=data,
search_space=search_space or experiment.search_space,
torch_dtype=dtype,
torch_device=device,
default_model_gen_options={
"acquisition_function_kwargs": {
"random_scalarization": True,
"sequential": True,
}
},
),
)
def get_MOO_PAREGO(
experiment: Experiment,
data: Data,
ref_point: Optional[List[float]] = None,
search_space: Optional[SearchSpace] = None,
dtype: torch.dtype = torch.double,
device: torch.device = DEFAULT_TORCH_DEVICE,
) -> TorchModelBridge:
"""Instantiates a multi-objective model that generates points with ParEGO.
qParEGO optimizes random augmented chebyshev scalarizations of the multiple
objectives. This allows it to explore non-convex pareto frontiers.
"""
# pyre-ignore: [16] `Optional` has no attribute `objective`.
if not isinstance(experiment.optimization_config.objective,
MultiObjective):
raise ValueError(
"Multi-Objective optimization requires multiple objectives")
if data.df.empty:
raise ValueError("MultiObjectiveOptimization requires non-empty data.")
return checked_cast(
TorchModelBridge,
Models.BOTORCH(
experiment=experiment,
data=data,
search_space=search_space or experiment.search_space,
torch_dtype=dtype,
torch_device=device,
ref_point=None,
default_model_gen_options={
"acquisition_function_kwargs": {
"chebyshev_scalarization": True,
"sequential": True,
}
},
),
)
def bayes_opt_loop(experiment,
model_constructor,
acquisition_function,
evaluation_budget,
n_initial_evaluations=0,
batch_size=1,
visualise=False,
vis_start=None,
vis_end=None,
vis_density=None,
function_name=None,
kernel_name=None,
log_dir=Path('./'),
observer=None):
print(f'Running Bayesian Optimisation')
# acqf_constructor = get_acqf_constructor(acquisition_function)
for i in range(n_initial_evaluations, evaluation_budget):
print(f'BO: {i + 1}/{evaluation_budget}')
model = Models.BOTORCH(
experiment=experiment,
data=experiment.eval(),
search_space=experiment.search_space,
model_constructor=model_constructor,
# acqf_constructor=acqf_constructor,
)
# Work out regrets, and add to observer
if observer is not None:
loc_regret, fun_regret = get_current_regrets(experiment, observer)
observer.current['results'][-1]['regret']['location'].append(
loc_regret.item())
observer.current['results'][-1]['regret']['function'].append(
fun_regret.item())
if visualise:
# go.Figure(
# data=plot_slice(model, 'x', 'function').data
# ).write_image(str(log_dir_i / 'posterior.png'))
mod_vis.posterior(model,
experiment,
vis_start,
vis_end,
vis_density,
function_name=function_name,
kernel_name=kernel_name,
log_dir=log_dir,
point_num=i)
mod_vis.kernel(model,
vis_start,
vis_end,
vis_density,
function_name=function_name,
kernel_name=kernel_name,
log_dir=log_dir,
point_num=i)
bo_vis.acquisition_function(model,
experiment,
vis_start,
vis_end,
vis_density,
log_dir=log_dir,
point_num=i)
vis.board(observer.current['function'],
model,
experiment,
observer.current['results'][-1]['regret']['location'],
observer.current['results'][-1]['regret']['function'],
n_initial_evaluations,
vis_start,
vis_end,
vis_density,
function_name=function_name,
kernel_name=kernel_name,
log_dir=log_dir,
point_num=i)
# Generate new point.
experiment.new_trial(generator_run=model.gen(batch_size))
def main():
args = parse_args()
function_list = [get_function_by_name[name]
for name in args.function_name_list.split(',')]
weight_list = list(map(float, args.weight_list.split(',')))
covariance_matrix = json.loads(args.covariance_matrix)
evaluate_covariance = args.evaluate_covariance
init_iter = args.init_iter
# if init_iter > 1:
# raise ValueError("init_iter should be 1.")
init_batch_size = args.init_batch_size
update_iter = args.update_iter
batch_size = args.batch_size
var_coef = args.var_coef
var_compute_type = args.var_compute_type
num_random = args.num_random
num_bucket = args.num_bucket
save_path = args.save_path
# num_control = args.num_control
minimize = True
groundtruth_function = get_groundtruth_function(function_list, weight_list)
#evaluation_function = get_evaluation_function(
# function_list, weight_list, covariance_matrix,
# evaluate_covariance, var_coef)
evaluation_function = get_evaluation_function(
function_list, weight_list, covariance_matrix,
var_compute_type, num_random, num_bucket)
exp = SimpleExperiment(
name=args.function_name_list + args.weight_list,
search_space=get_search_space(function_list),
evaluation_function=evaluation_function,
objective_name="objective_name",
minimize=minimize,
)
t_start = time.time()
print(f"Start time: {t_start}")
print(f"Sobol iteration begin...{time.time() - t_start}")
sobol = Models.SOBOL(exp.search_space)
for i in range(init_iter):
if init_batch_size == 1:
exp.new_trial(generator_run=sobol.gen(init_batch_size))
else:
exp.new_batch_trial(generator_run=sobol.gen(init_batch_size))
print(f"Running sobol optimization trial {i+1}/{init_iter}..."
f"{time.time() - t_start}")
print(f"GPEI iteration begin...{time.time() - t_start}")
for i in range(update_iter):
gpei = Models.BOTORCH(experiment=exp, data=exp.eval())
if batch_size == 1:
exp.new_trial(generator_run=gpei.gen(batch_size))
else:
exp.new_batch_trial(generator_run=gpei.gen(batch_size))
print(f"Running GPEI optimization trial {i+1}/{update_iter}..."
f"{time.time() - t_start}")
# Construct Result.
## origin data.
data_df = copy.deepcopy(exp.eval().df)
compare_func = min if minimize else max
arm_name2mean = {}
for _, row in data_df.iterrows():
arm_name2mean[row["arm_name"]] = row["mean"]
## parameters true_mean.
other_columns = {
"arm_name": [], "parameters": [], "true_mean": [],
"cur_trial_best_mean": [], "accum_trials_best_mean": []}
atbm = None # accum_trial_best_mean
for trial in exp.trials.values():
ctbm = None # cur_trial_best_mean
for arm in trial.arms:
other_columns['arm_name'].append(arm.name)
other_columns['parameters'].append(json.dumps(arm.parameters))
other_columns['true_mean'].append(
groundtruth_function(arm.parameters))
if ctbm is None:
ctbm = arm_name2mean[arm.name]
ctbm = compare_func(ctbm, arm_name2mean[arm.name])
if atbm is None:
atbm = ctbm
atbm = compare_func(atbm, ctbm)
other_columns['cur_trial_best_mean'].extend([ctbm] * len(trial.arms))
other_columns['accum_trials_best_mean'].extend([atbm] * len(trial.arms))
other_df = DataFrame(other_columns)
result_df = data_df.set_index('arm_name').join(
other_df.set_index('arm_name')).reset_index()
# Save to file.
print("Save to file.")
sub_dir_name = "_".join([
"ax", args.function_name_list.replace(",", "_"),
args.weight_list.replace(",", "_"), args.covariance_matrix.replace(
"[", "_").replace("]", "_").replace(",", "_").replace(" ", ""),
str(args.evaluate_covariance), str(args.init_iter), str(init_batch_size),
str(args.update_iter), str(args.batch_size), str(args.var_coef),
str(minimize), str(var_compute_type), str(num_random), str(num_bucket)
])
abs_dir_path = os.path.join(save_path, sub_dir_name)
Path(abs_dir_path).mkdir(parents=True, exist_ok=True)
task_id = os.environ.get('TASK_INDEX')
cur_time = pd.Timestamp.now().strftime('%Y%m%d%H%M%S')
filename = cur_time + "_" + str(task_id) + ".csv"
print(os.path.join(abs_dir_path, filename))
result_df.to_csv(os.path.join(abs_dir_path, filename))
print("2021-01-19 19:48:00")
print("Done...")
def test_get_standard_plots(self):
exp = get_branin_experiment()
self.assertEqual(
len(
get_standard_plots(experiment=exp,
model=get_generation_strategy().model)),
0,
)
exp = get_branin_experiment(with_batch=True, minimize=True)
exp.trials[0].run()
plots = get_standard_plots(
experiment=exp,
model=Models.BOTORCH(experiment=exp, data=exp.fetch_data()),
)
self.assertEqual(len(plots), 6)
self.assertTrue(all(isinstance(plot, go.Figure) for plot in plots))
exp = get_branin_experiment_with_multi_objective(with_batch=True)
exp.optimization_config.objective.objectives[0].minimize = False
exp.optimization_config.objective.objectives[1].minimize = True
exp.optimization_config._objective_thresholds = [
ObjectiveThreshold(metric=exp.metrics["branin_a"],
op=ComparisonOp.GEQ,
bound=-100.0),
ObjectiveThreshold(metric=exp.metrics["branin_b"],
op=ComparisonOp.LEQ,
bound=100.0),
]
exp.trials[0].run()
plots = get_standard_plots(experiment=exp,
model=Models.MOO(experiment=exp,
data=exp.fetch_data()))
self.assertEqual(len(plots), 7)
# All plots are successfully created when objective thresholds are absent
exp.optimization_config._objective_thresholds = []
plots = get_standard_plots(experiment=exp,
model=Models.MOO(experiment=exp,
data=exp.fetch_data()))
self.assertEqual(len(plots), 7)
exp = get_branin_experiment_with_timestamp_map_metric(
with_status_quo=True)
exp.new_trial().add_arm(exp.status_quo)
exp.trials[0].run()
exp.new_trial(generator_run=Models.SOBOL(
search_space=exp.search_space).gen(n=1))
exp.trials[1].run()
plots = get_standard_plots(
experiment=exp,
model=Models.BOTORCH(experiment=exp, data=exp.fetch_data()),
true_objective_metric_name="branin",
)
self.assertEqual(len(plots), 9)
self.assertTrue(all(isinstance(plot, go.Figure) for plot in plots))
self.assertIn(
"Objective branin_map vs. True Objective Metric branin",
[p.layout.title.text for p in plots],
)
with self.assertRaisesRegex(
ValueError, "Please add a valid true_objective_metric_name"):
plots = get_standard_plots(
experiment=exp,
model=Models.BOTORCH(experiment=exp, data=exp.fetch_data()),
true_objective_metric_name="not_present",
)
def test_best_from_model_prediction(self):
exp = get_branin_experiment()
for _ in range(3):
sobol = Models.SOBOL(search_space=exp.search_space)
generator_run = sobol.gen(n=1)
trial = exp.new_trial(generator_run=generator_run)
trial.run()
trial.mark_completed()
exp.attach_data(exp.fetch_data())
gpei = Models.BOTORCH(experiment=exp, data=exp.lookup_data())
generator_run = gpei.gen(n=1)
trial = exp.new_trial(generator_run=generator_run)
trial.run()
trial.mark_completed()
with patch.object(
ArrayModelBridge,
"model_best_point",
return_value=(
(
Arm(
name="0_0",
parameters={"x1": -4.842811906710267, "x2": 11.887089014053345},
),
(
{"branin": 34.76260622783635},
{"branin": {"branin": 0.00028306433439807734}},
),
)
),
) as mock_model_best_point, self.assertLogs(
logger="ax.service.utils.best_point", level="WARN"
) as lg:
# Test bad model fit causes function to resort back to raw data
with patch(
"ax.service.utils.best_point.assess_model_fit",
return_value=AssessModelFitResult(
good_fit_metrics_to_fisher_score={},
bad_fit_metrics_to_fisher_score={
"branin": 0,
},
),
):
self.assertIsNotNone(get_best_parameters(exp, Models))
self.assertTrue(
any("Model fit is poor" in warning for warning in lg.output),
msg=lg.output,
)
mock_model_best_point.assert_not_called()
# Test model best point is used when fit is good
with patch(
"ax.service.utils.best_point.assess_model_fit",
return_value=AssessModelFitResult(
good_fit_metrics_to_fisher_score={
"branin": 0,
},
bad_fit_metrics_to_fisher_score={},
),
):
self.assertIsNotNone(get_best_parameters(exp, Models))
mock_model_best_point.assert_called()
# Assert the non-mocked method works correctly as well
self.assertIsNotNone(get_best_parameters(exp, Models))
print(f" - {len(ran)} Parameters {p}")
else:
continue
# ax_client.attach_trial(parameters=p)
for t in exp.trials:
exp.trials[t].run()
# Pull out of API for more specific iteration
batch = 10
print(f"Running Batch GP+EI optimization of {batch} samples")
# Reinitialize GP+EI model at each step with updated data.
data = exp.fetch_data()
print(data.df)
gpei = Models.BOTORCH(experiment=exp, data=data)
generator = gpei.gen(batch)
exp.new_batch_trial(generator_run=generator)
new_trial = len(exp.trials) - 1
print(f"New Candidate Designs")
for arm in exp.trials[new_trial].arms:
print(arm)
# generator_run = gpei.gen(5)
# experiment.new_batch_trial(generator_run=generator_run)
#
# quit()
# sobol = Models.SOBOL(search_space=experiment.search_space)
# generator_run = sobol.gen(5)
def pid(cfg):
env_name = cfg.env.params.name
env = gym.make(env_name)
env.reset()
full_rewards = []
exp_cfg = cfg.experiment
from learn.utils.plotly import hv_characterization
hv_characterization()
def compare_control(env, cfg, save=True):
import torch
from learn.control.pid import PidPolicy
controllers = []
labels = []
metrics = []
# PID baselines
# /Users/nato/Documents/Berkeley/Research/Codebases/dynamics-learn/sweeps/2020-04-14/11-12-02
# from learn.simulate_sac import *
# Rotation policy
sac_policy1 = torch.load(
'/Users/nato/Documents/Berkeley/Research/Codebases/dynamics-learn/outputs/2020-03-24/18-32-26/trial_70000.dat'
)
controllers.append(sac_policy1['policy'])
labels.append("SAC - Rotation")
metrics.append(0)
# Living reward policy
sac_policy2 = torch.load(
'/Users/nato/Documents/Berkeley/Research/Codebases/dynamics-learn/outputs/2020-03-24/18-31-45/trial_35000.dat'
)
controllers.append(sac_policy2['policy'])
labels.append("SAC - Living")
metrics.append(1)
# Square cost policy
# sac_policy2 = torch.load(
# '/Users/nato/Documents/Berkeley/Research/Codebases/dynamics-learn/sweeps/2020-03-25/20-30-47/metric.name=Square,robot=iono_sim/26/trial_40000.dat')
controllers.append(sac_policy2['policy'])
labels.append("SAC - Square")
metrics.append(2)
# un-Optimized PID parameters
pid_params = [[2531.917, 61.358, 33.762], [2531.917, 61.358, 33.762]]
pid = PidPolicy(cfg)
pid.set_params(pid_params)
controllers.append(pid)
labels.append("PID - temp")
metrics.append(0)
controllers.append(pid)
labels.append("PID - temp")
metrics.append(1)
# Optimized PID parameters
pid_params = [[2531.917, 61.358, 3333.762],
[2531.917, 61.358, 3333.762]]
pid = PidPolicy(cfg)
pid.set_params(pid_params)
controllers.append(pid)
labels.append("PID - improved")
metrics.append(2)
from learn.control.mpc import MPController
cfg.policy.mode = 'mpc'
# dynam_model = torch.load(
# '/Users/nato/Documents/Berkeley/Research/Codebases/dynamics-learn/outputs/2020-03-25/10-45-17/trial_1.dat')
dynam_model = torch.load(
'/Users/nato/Documents/Berkeley/Research/Codebases/dynamics-learn/sweeps/2020-03-25/20-30-57/metric.name=Rotation,robot=iono_sim/14/trial_9.dat'
)
mpc = MPController(env, dynam_model['model'], cfg)
controllers.append(mpc)
labels.append("MPC - 1")
metrics.append(0)
controllers.append(mpc)
labels.append("MPC - 2")
metrics.append(1)
controllers.append(mpc)
labels.append("MPC - 3")
metrics.append(2)
import plotly.graph_objects as go
import plotly
colors = [
'#1f77b4', # muted blue
'#ff7f0e', # safety orange
'#2ca02c', # cooked asparagus green
'#d62728', # brick red
'#9467bd', # muted purple
'#8c564b', # chestnut brown
'#e377c2', # raspberry yogurt pink
'#7f7f7f', # middle gray
'#bcbd22', # curry yellow-green
'#17becf' # blue-teal
]
markers = [
"cross",
"circle-open-dot",
"x-open-dot",
"triangle-up-open-dot",
"y-down-open",
"diamond-open-dot",
"hourglass",
"hash",
"star",
"square",
]
m1 = living_reward
m2 = rotation_mat
m3 = squ_cost
eval_metrics = [m1, m2, m3]
metric_names = ["Living", "Rotation", "Square"]
fig = plotly.subplots.make_subplots(
rows=3,
cols=2,
# subplot_titles=["Living", "Rotation", "Square"],
subplot_titles=[
"Pitch",
"Roll",
" ",
" ",
" ",
" ",
],
vertical_spacing=0.03,
horizontal_spacing=0.03,
shared_xaxes=True,
) # go.Figure()
fig_mpc = go.Figure()
fig_sac = go.Figure()
pry = [1, 0, 2]
# state0 = 2*env.reset()
# state0 = env.reset()
state0 = np.array([0, np.deg2rad(15), 0, 0, 0, 0])
for i, (con, lab, m) in enumerate(zip(controllers, labels, metrics)):
print(f"Evaluating controller type {lab}")
_ = env.reset()
env.set_state(np.concatenate((np.zeros(6), state0)))
state = state0
states = []
actions = []
rews = []
done = False
# for t in range(cfg.experiment.r_len + 1):
for t in range(500):
if done:
break
if "SAC" in lab:
with torch.no_grad():
with eval_mode(con):
action = con.select_action(state)
if i < 2:
action = np.array([65535, 65535, 65535, 65535
]) * (action + 1) / 2
else:
action = np.array([3000, 3000, 3000, 3000
]) * (action + 1) / 2
else:
action = con.get_action(state, metric=eval_metrics[m])
states.append(state)
actions.append(action)
state, rew, done, _ = env.step(action)
done = done
states = np.stack(states)
actions = np.stack(actions)
pitch = np.degrees(states[:, pry[0]])
roll = np.degrees(states[:, pry[1]])
# deal with markers
num_mark = np.zeros(len(pitch))
mark_every = 50
m_size = 32
start = np.random.randint(0, int(len(pitch) / 10))
num_mark[start::mark_every] = m_size
if "SAC" in lab:
fig_sac.add_trace(
go.Scatter(
y=pitch,
name=metric_names[m], # legendgroup=lab[:3],
# showlegend=(True if (i % 3 == 0) else False),
line=dict(color=colors[m], width=4),
cliponaxis=False,
mode='lines+markers',
marker=dict(color=colors[m],
symbol=markers[-m],
size=num_mark.tolist())))
elif "MPC" in lab:
fig_mpc.add_trace(
go.Scatter(
y=pitch,
name=metric_names[m], # legendgroup=lab[:3],
# showlegend=(True if (i % 3 == 0) else False),
line=dict(color=colors[m], width=4),
cliponaxis=False,
mode='lines+markers',
marker=dict(color=colors[m],
symbol=markers[-m],
size=num_mark.tolist())))
fig.add_trace(
go.Scatter(
y=pitch,
name=lab[:3] + str(m),
legendgroup=lab[:3],
showlegend=(True if (i % 3 == 0) else False),
line=dict(color=colors[int(i / 3)],
width=2), # mode='lines+markers',
# marker=dict(color=colors[i], symbol=markers[i], size=16)
),
row=m + 1,
col=1)
fig.add_trace(
go.Scatter(
y=roll,
name=lab[:3] + str(m),
legendgroup=lab[:3],
showlegend=(False),
line=dict(color=colors[int(i / 3)],
width=2), # mode='lines+markers',
# marker=dict(color=colors[i], symbol=markers[i], size=16)
),
row=m + 1,
col=2)
fig.update_layout(
title='Comparison of Controllers and Reward Functions',
font=dict(family="Times New Roman, Times, serif",
size=24,
color="black"),
legend_orientation="h",
legend=dict(
x=.6,
y=0.07,
bgcolor='rgba(205, 223, 212, .4)',
bordercolor="Black",
),
# xaxis_title='Timestep',
# yaxis_title='Angle (Degrees)',
plot_bgcolor='white',
width=1600,
height=1000,
# xaxis=dict(
# showline=True,
# showgrid=False,
# showticklabels=True, ),
# yaxis=dict(
# showline=True,
# showgrid=False,
# showticklabels=True, ),
)
fig_sac.update_layout( # title='Comparison of SAC Policies',
font=dict(family="Times New Roman, Times, serif",
size=32,
color="black"),
legend_orientation="h",
legend=dict(
x=.35,
y=0.1,
bgcolor='rgba(205, 223, 212, .4)',
bordercolor="Black",
),
# xaxis_title='Timestep',
# yaxis_title='Angle (Degrees)',
showlegend=False,
plot_bgcolor='white',
width=1600,
height=800,
margin=dict(t=5, r=5),
)
fig_mpc.update_layout( # title='Comparison of MPC Policies',
font=dict(family="Times New Roman, Times, serif",
size=32,
color="black"),
legend_orientation="h",
showlegend=False,
legend=dict(
x=.35,
y=0.1,
bgcolor='rgba(205, 223, 212, .4)',
bordercolor="Black",
# ncol= 2,
),
# xaxis_title='Timestep',
# yaxis_title='Angle (Degrees)',
plot_bgcolor='white',
width=1600,
height=800,
margin=dict(t=5, r=5),
)
reg_color = 'rgba(255,60,60,.15)'
fig_sac.add_trace(
go.Scatter(x=[0, 500],
y=[5, 5],
name='Living Region',
legendgroup='Living Region',
fill='tozeroy',
mode='lines',
fillcolor=reg_color,
line=dict(width=0.0,
color=reg_color))) # fill down to xaxis
fig_sac.add_trace(
go.Scatter(x=[0, 500],
y=[-5, -5],
showlegend=False,
legendgroup='Living Region',
fill='tozeroy',
mode='lines',
fillcolor=reg_color,
line=dict(width=0.0,
color=reg_color))) # fill down to xaxis
fig_mpc.add_trace(
go.Scatter(x=[0, 500],
y=[5, 5],
name='Living Region',
legendgroup='Living Region',
fill='tozeroy',
mode='lines',
fillcolor=reg_color,
line=dict(width=0.0,
color=reg_color))) # fill down to xaxis
fig_mpc.add_trace(
go.Scatter(x=[0, 500],
y=[-5, -5],
showlegend=False,
legendgroup='Living Region',
fill='tozeroy',
mode='lines',
fillcolor=reg_color,
line=dict(width=0.0,
color=reg_color))) # fill down to xaxis
# SOLO
rang_ind = [-20, 20]
fig_sac.update_xaxes(
title_text="Timestep",
range=[0, 500],
ticks="inside",
tickwidth=2,
zeroline=True,
zerolinecolor='rgba(0,0,0,.5)',
zerolinewidth=1,
)
fig_sac.update_yaxes(
title_text="Pitch (degrees)",
range=rang_ind,
ticks="inside",
tickwidth=2,
zeroline=True,
zerolinecolor='rgba(0,0,0,.5)',
zerolinewidth=1,
)
fig_sac.show()
fig_sac.write_image(os.getcwd() + "/compare_sac.pdf")
fig_mpc.update_xaxes(
title_text="Timestep",
range=[0, 500],
ticks="inside",
tickwidth=2,
zeroline=True,
zerolinecolor='rgba(0,0,0,.5)',
zerolinewidth=1,
)
fig_mpc.update_yaxes(
title_text="Pitch (degrees)",
range=rang_ind,
ticks="inside",
tickwidth=2,
zeroline=True,
zerolinecolor='rgba(0,0,0,.5)',
zerolinewidth=1,
)
fig_mpc.show()
fig_mpc.write_image(os.getcwd() + "/compare_mpc.pdf")
# COMPARISON
fig.update_xaxes(
title_text="Timestep",
row=3,
col=1,
ticks="inside",
tickwidth=2,
zeroline=True,
zerolinecolor='rgba(0,0,0,.5)',
zerolinewidth=1,
)
fig.update_xaxes(
row=2,
col=1,
zeroline=True,
zerolinecolor='rgba(0,0,0,.5)',
zerolinewidth=1,
)
fig.update_xaxes(
row=1,
col=1,
zeroline=True,
zerolinecolor='rgba(0,0,0,.5)',
zerolinewidth=1,
)
fig.update_xaxes(
title_text="Timestep",
row=3,
col=2,
ticks="inside",
tickwidth=2,
zeroline=True,
zerolinecolor='rgba(0,0,0,.5)',
zerolinewidth=1,
)
fig.update_xaxes(
row=2,
col=2,
zeroline=True,
zerolinecolor='rgba(0,0,0,.5)',
zerolinewidth=1,
)
fig.update_xaxes(
row=1,
col=2,
zeroline=True,
zerolinecolor='rgba(0,0,0,.5)',
zerolinewidth=1,
)
# fig.update_xaxes(title_text="xaxis 1 title", row=1, col=1)
# fig.update_yaxes(title_text="Roll (Degrees)", row=1, col=1)
rang = [-30, 30]
nticks = 6
fig.update_yaxes(
title_text="Living Rew.",
range=rang,
row=1,
col=1,
nticks=nticks,
ticks="inside",
tickwidth=2,
zeroline=True,
zerolinecolor='rgba(0,0,0,.5)',
zerolinewidth=1,
)
fig.update_yaxes(
title_text="Rotation Rew.",
range=rang,
row=2,
col=1,
nticks=nticks,
ticks="inside",
tickwidth=2,
zeroline=True,
zerolinecolor='rgba(0,0,0,.5)',
zerolinewidth=1,
)
fig.update_yaxes(
title_text="Square Cost",
range=rang,
row=3,
col=1,
nticks=nticks,
ticks="inside",
tickwidth=2,
zeroline=True,
zerolinecolor='rgba(0,0,0,.5)',
zerolinewidth=1,
)
fig.update_yaxes(
range=rang,
row=1,
col=2,
nticks=nticks,
showticklabels=False,
ticks="inside",
tickwidth=2,
zeroline=True,
zerolinecolor='rgba(0,0,0,.5)',
zerolinewidth=1,
)
fig.update_yaxes(
range=rang,
row=2,
col=2,
nticks=nticks,
showticklabels=False,
ticks="inside",
tickwidth=2,
zeroline=True,
zerolinecolor='rgba(0,0,0,.5)',
zerolinewidth=1,
)
fig.update_yaxes(
range=rang,
row=3,
col=2,
nticks=nticks,
showticklabels=False,
ticks="inside",
tickwidth=2,
zeroline=True,
zerolinecolor='rgba(0,0,0,.5)',
zerolinewidth=1,
)
print(f"Plotting {len(labels)} control responses")
# save = False
# if save:
# fig.write_image(os.getcwd() + "compare.png")
# else:
# fig.show()
#
# return fig
# compare_control(env, cfg, save=True)
# quit()
plot_results(logx=False, save=True, mpc=False)
quit()
# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
# # # # # # # # # # # Evalutation Function # # # # # # # # # # # # # # # # # # # #
def bo_rollout_wrapper(params, weights=None): # env, controller, exp_cfg):
pid_1 = [params["pitch-p"], params["pitch-i"], params["pitch-d"]]
# pid_1 = [params["roll-p"], params["roll-i"],
# params["roll-d"]] # [params["pitch-p"], params["pitch-i"], params["pitch-d"]]
pid_2 = [params["roll-p"], params["roll-i"], params["roll-d"]]
print(
f"Optimizing Parameters {np.round(pid_1, 3)},{np.round(pid_2, 3)}")
pid_params = [[pid_1[0], pid_1[1], pid_1[2]],
[pid_2[0], pid_2[1], pid_2[2]]]
# pid_params = [[1000, 0, 0], [1000, 0, 0]]
pid = PidPolicy(cfg)
pid.set_params(pid_params)
cum_cost = []
r = 0
fncs = [squ_cost, living_reward, rotation_mat]
mult_rewards = [[] for _ in range(len(fncs))]
while r < cfg.experiment.repeat:
pid.reset()
states, actions, rews, sim_error = rollout(env, pid, exp_cfg)
# plot_rollout(states, actions, pry=[1, 0, 2])
rewards_full = get_rewards(states, actions, fncs=fncs)
for i, vec in enumerate(rewards_full):
mult_rewards[i].append(vec)
# if sim_error:
# print("Repeating strange simulation")
# continue
# if len(rews) < 400:
# cum_cost.append(-(cfg.experiment.r_len - len(rews)) / cfg.experiment.r_len)
# else:
rollout_cost = np.sum(rews) / cfg.experiment.r_len # / len(rews)
# if rollout_cost > max_cost:
# max_cost = rollout_cost
# rollout_cost += get_reward_euler(states[-1], actions[-1])
cum_cost.append(rollout_cost)
r += 1
std = np.std(cum_cost)
cum_cost = np.mean(cum_cost)
# print(f"Cum. Cost {cum_cost / max_cost}")
# print(f"- Mean Cum. Cost / Rew: {cum_cost}, std dev: {std}")
eval = {
"Square": (np.mean(rewards_full[0]), np.std(rewards_full[0])),
"Living": (np.mean(rewards_full[1]), np.std(rewards_full[1])),
"Rotation": (np.mean(rewards_full[2]), np.std(rewards_full[2]))
}
for n, (key, value) in enumerate(eval.items()):
if n == 0:
print(f"- Square {np.round(value, 4)}")
elif n == 1:
print(f"- Living {np.round(value, 4)}")
else:
print(f"- Rotn {np.round(value, 4)}")
return eval
# return cum_cost.reshape(1, 1), std
# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
from ax import (
ComparisonOp,
ParameterType,
RangeParameter,
SearchSpace,
SimpleExperiment,
OutcomeConstraint,
)
exp = SimpleExperiment(
name="PID Control Robot",
search_space=SearchSpace([
RangeParameter(
name=f"roll-p",
parameter_type=ParameterType.FLOAT,
lower=1.0,
upper=10000.0,
log_scale=True,
),
# FixedParameter(name="roll-i", value=0.0, parameter_type=ParameterType.FLOAT),
RangeParameter(
name=f"roll-i",
parameter_type=ParameterType.FLOAT,
lower=0,
upper=1000.0,
log_scale=False,
),
RangeParameter(
name=f"roll-d",
parameter_type=ParameterType.FLOAT,
lower=.1,
upper=5000.0,
log_scale=True,
),
RangeParameter(
name=f"pitch-p",
parameter_type=ParameterType.FLOAT,
lower=1.0,
upper=10000.0,
log_scale=True,
),
RangeParameter(
name=f"pitch-d",
parameter_type=ParameterType.FLOAT,
lower=0,
upper=1000.0,
log_scale=False,
),
RangeParameter(
name=f"pitch-i",
parameter_type=ParameterType.FLOAT,
lower=.1,
upper=5000.0,
log_scale=True,
),
# FixedParameter(name="pitch-i", value=0.0, parameter_type=ParameterType.FLOAT),
]),
evaluation_function=bo_rollout_wrapper,
objective_name=cfg.metric.name,
minimize=cfg.metric.minimize,
outcome_constraints=[],
)
from ax.storage.metric_registry import register_metric
from ax.storage.runner_registry import register_runner
class GenericMetric(Metric):
def fetch_trial_data(self, trial):
records = []
for arm_name, arm in trial.arms_by_name.items():
params = arm.parameters
mean, sem = bo_rollout_wrapper(params)
records.append({
"arm_name": arm_name,
"metric_name": self.name,
"mean": mean,
"sem": sem,
"trial_index": trial.index,
})
return Data(df=pd.DataFrame.from_records(records))
class MyRunner(Runner):
def run(self, trial):
return {"name": str(trial.index)}
optimization_config = OptimizationConfig(objective=Objective(
metric=GenericMetric(name=cfg.metric.name),
minimize=cfg.metric.minimize,
), )
register_metric(GenericMetric)
register_runner(MyRunner)
exp.runner = MyRunner()
exp.optimization_config = optimization_config
log.info(f"Running experiment, metric name {cfg.metric.name}")
log.info(f"Running Sobol initialization trials...")
sobol = Models.SOBOL(exp.search_space)
num_search = cfg.bo.random
for i in range(num_search):
exp.new_trial(generator_run=sobol.gen(1))
exp.trials[len(exp.trials) - 1].run()
import plotly.graph_objects as go
gpei = Models.BOTORCH(experiment=exp, data=exp.eval())
objectives = ["Living", "Square", "Rotation"]
def plot_all(model, objectives, name="", rend=False):
for o in objectives:
plot = plot_contour(
model=model,
param_x="roll-p",
param_y="roll-d",
metric_name=o,
)
plot[0]['layout']['title'] = o
data = plot[0]['data']
lay = plot[0]['layout']
for i, d in enumerate(data):
if i > 1:
d['cliponaxis'] = False
fig = {
"data": data,
"layout": lay,
}
go.Figure(fig).write_image(name + o + ".png")
if rend: render(plot)
plot_all(gpei, objectives, name="Random fit-")
num_opt = cfg.bo.optimized
for i in range(num_opt):
log.info(f"Running GP+EI optimization trial {i + 1}/{num_opt}...")
# Reinitialize GP+EI model at each step with updated data.
batch = exp.new_trial(generator_run=gpei.gen(1))
gpei = Models.BOTORCH(experiment=exp, data=exp.eval())
if ((i + 1) % 10) == 0:
plot_all(gpei,
objectives,
name=f"optimizing {str(i + 1)}-",
rend=False)
from ax.plot.exp_utils import exp_to_df
best_arm, _ = gpei.gen(1).best_arm_predictions
best_parameters = best_arm.parameters
log.info(f"Best parameters {best_parameters}")
experiment_log = {
"Exp": exp_to_df(exp=exp),
"Cfg": cfg,
"Best_param": best_parameters,
}
log.info("Printing Parameters")
log.info(exp_to_df(exp=exp))
save_log(cfg, exp, experiment_log)
fig_learn = plot_learning(exp, cfg)
fig_learn.write_image("learning" + ".png")
fig_learn.show()
plot_all(gpei, objectives, name=f"FINAL -", rend=True)