SAM.LegsOptimization/optimize.py at main · SofaComplianceRobotics/SAM.LegsOptimization · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
import sys
import os
import importlib
import argparse

#import param
import optuna
import optunahub
from optimization.callbacks import LogCompletedTrial, SaveCompletedTrial, RepeatTrial
from optimization.processing import logResults, createFolder, appendDataset
# Argument parser
parser = argparse.ArgumentParser(description='Optimize an objective function.')
parser.add_argument('objective',
                    metavar='module.py',
                    type=str,
                    nargs=1,
                    help='a python script that should contain an objective ' \
                         'function with the following signature: '
                         'def objective(trial). See objective.py in this directory for example')

parser.add_argument('-n',
                    '--nTrials',
                    type=int, nargs='?',
                    help='number of trials (default: 20)',
                    default=20)

args = parser.parse_args()

# Import the objective module
print("Importing objective module from: ", args.objective[0])
module = args.objective[0]
sys.path.append(os.path.dirname(os.path.abspath(module)))
sys.path.append(os.path.dirname(os.path.abspath(sys.argv[0])))
sys.path.append(os.path.dirname(os.path.abspath(sys.argv[0]))+"/../..")
o = importlib.import_module(os.path.splitext(os.path.basename(module))[0])

# Create a study for optimization

storage = optuna.storages.InMemoryStorage()

#sampler = optuna.samplers.CmaEsSampler()

#sampler = optuna.samplers.TPESampler()

# module = optunahub.load_module(package="samplers/hill_climbing")
# sampler = module.HillClimbingSampler()

# module = optunahub.load_module(package="samplers/auto_sampler")
# sampler = module.AutoSampler()

# module = optunahub.load_module(package="samplers/nelder_mead")
# sampler = module.NelderMeadSampler(seed=123)

module = optunahub.load_module("samplers/cma_es_refinement")
sampler = module.CmaEsRefinementSampler(seed=42)

pruner = optuna.pruners.HyperbandPruner()
#pruner = optuna.pruners.ThresholdPruner(upper=5)

study = optuna.create_study(storage = storage, sampler = sampler, pruner =pruner)

# Optimize the objective function
print(o.objective.__doc__)

study.optimize(func=o.objective,
               gc_after_trial=True, #Auto garbage clean to not saturate the memory
               n_trials=args.nTrials,
               #n_jobs=5,
               show_progress_bar=True,
               callbacks=[SaveCompletedTrial(study_name=study.study_name),
                          #RepeatTrial(maxIteration=5),
                          #optuna.study.MaxTrialsCallback(n_trials=5, states=(optuna.trial.TrialState.COMPLETE,)),
                          #LogCompletedTrial()
                          ]
               )

# Log optimization results
logResults(study)
appendDataset(study,'nDataset.csv')

# Visualize the optimization history
fig = optuna.visualization.plot_optimization_history(study)
fig.show()