Parameter tuning of MPC via Optuna

optuna_mpc.py

from casadi import *
import math
import numpy as np
import optuna


class CartPole:
    def __init__(self):
        self.mc = 2.0   # cart の質量[kg]
        self.mp = 0.2   # pole の質量[kg]
        self.l = 0.5    # pole の長さ[m]
        self.ga = 9.81   # 重力加速度[m/s2]

    def dynamics(self, x, u):
        mc = self.mc
        mp = self.mp
        l = self.l
        ga = self.ga
        y = x[0]     # cart の水平位置[m]
        th = x[1]    # pole の傾き角[rad]
        dy = x[2]    # cart の水平速度[m/s]
        dth = x[3]   # pole の傾き角速度[rad/s]
        f = u     # cart を押す力[N]（制御入力）
        # cart の水平加速度
        ddy = (f+mp*sin(th)*(l*dth*dth+ga*cos(th))) / (mc+mp*sin(th)*sin(th)) 
        # pole の傾き角加速度
        ddth = (-f*cos(th)-mp*l*dth*dth*cos(th)*sin(th)-(mc+mp)*ga*sin(th)) / (l * (mc+mp*sin(th)*sin(th))) 
        # Forward Euler による離散化状態方程式
        return dy, dth, ddy, ddth


class TrainParams:
    def __init__(self, T=1.0, N=20, dy_weight=0.01, dth_weight=0.01, u_weight=0.1, max_ipopt_iter=5):
        self.T = T
        self.N = N
        self.dy_weight = dy_weight
        self.dth_weight = dth_weight
        self.u_weight = u_weight
        self.max_ipopt_iter = max_ipopt_iter


class MPC:
    def __init__(self, train_params):
        # 学習パラメータ
        T = train_params.T
        N = train_params.N
        dy_weight = train_params.dy_weight
        dth_weight = train_params.dth_weight
        u_weight = train_params.u_weight
        max_ipopt_iter = train_params.max_ipopt_iter
        # 問題設定
        cartpole = CartPole()  # ダイナミクスモデル
        dt = T / N  # 離散化ステップ
        nx = 4      # 状態空間の次元
        nu = 1      # 制御入力の次元

        # 以下で非線形計画問題(NLP)を定式化
        w = []    # 最適化変数を格納する list
        w0 = []   # 最適化変数(w)の初期推定解を格納する list
        lbw = []  # 最適化変数(w)の lower bound を格納する list
        ubw = []  # 最適化変数(w)の upper bound を格納する list
        J = 0     # コスト関数 
        g = []    # 制約（等式制約，不等式制約どちらも）を格納する list
        lbg = []  # 制約関数(g)の lower bound を格納する list
        ubg = []  # 制約関数(g)の upper bound を格納する list
        lam_x0 = []  # 制約 lbw<w<ubw のラグランジュ乗数
        lam_g0 = []  # 制約 lbg<g<ubg のラグランジュ乗数

        Xk = MX.sym('X0', nx) # 初期時刻の状態ベクトル x0
        w += [Xk]             # x0 を 最適化変数 list (w) に追加
        # 初期状態は given という条件を等式制約として考慮
        lbw += [0, 0, 0, 0]  # 等式制約は lower-bound と upper-bound を同じ値にすることで設定
        ubw += [0, 0, 0, 0]      # 等式制約は lower-bound と upper-bound を同じ値にすることで設定
        w0 +=  [0, 0, 0, 0]      # x0 の初期推定解
        lam_x0 += [0, 0, 0, 0]    # ラグランジュ乗数の初期推定解

        # 離散化ステージ 0~N-1 までのコストと制約を設定
        for k in range(N):
            Uk = MX.sym('U_' + str(k), nu) # 時間ステージ k の制御入力 uk を表す変数
            w   += [Uk]                   # uk を最適化変数 list に追加
            lbw += [-15.0]                # uk の lower-bound
            ubw += [15.0]                 # uk の upper-bound
            w0  += [0]                    # uk の初期推定解
            lam_x0 += [0]                 # ラグランジュ乗数の初期推定解

            # ステージコスト
            y = Xk[0]     # cart の水平位置[m]
            th = Xk[1]    # pole の傾き角[rad]
            dy = Xk[2]    # cart の水平速度[m/s]
            dth = Xk[3]   # pole の傾き角速度[rad/s]
            f = Uk[0]     # cart を押す力[N]（制御入力）
            # ステージコストのパラメータ 
            x_ref = [0.0, math.pi, 0.0, 0.0]  # 目標状態
            Q = [1.0, 10.0, dy_weight, dth_weight] # 状態への重み (学習パラメータ)
            R = [u_weight]                    # 制御入力への重み (学習パラメータ)
            L = 0                             # ステージコスト
            for i in range(nx):
                L += 0.5 * Q[i] * (Xk[i]-x_ref[i])**2 
            for i in range(nu):
                L += 0.5 * R[i] * Uk[i]**2
            J = J + L * dt                    # コスト関数にステージコストを追加

            # Forward Euler による離散化状態方程式
            dy, dth, ddy, ddth = cartpole.dynamics(Xk, Uk)
            Xk_next = vertcat(y + dy * dt, 
                              th + dth * dt,
                              dy + ddy * dt,
                              dth + ddth * dt)
            Xk1 = MX.sym('X_' + str(k+1), nx)  # 時間ステージ k+1 の状態 xk+1 を表す変数
            w   += [Xk1]                       # xk+1 を最適化変数 list に追加
            lbw += [-inf, -inf, -inf, -inf]    # xk+1 の lower-bound （指定しない要素は -inf）
            ubw += [inf, inf, inf, inf]        # xk+1 の upper-bound （指定しない要素は inf）
            w0 += [0.0, 0.0, 0.0, 0.0]         # xk+1 の初期推定解
            lam_x0 += [0, 0, 0, 0]              # ラグランジュ乗数の初期推定解

            # 状態方程式(xk+1=xk+fk*dt) を等式制約として導入
            g   += [Xk_next-Xk1]
            lbg += [0, 0, 0, 0]     # 等式制約は lower-bound と upper-bound を同じ値にすることで設定
            ubg += [0, 0, 0, 0]     # 等式制約は lower-bound と upper-bound を同じ値にすることで設定
            lam_g0 += [0, 0, 0, 0]   # ラグランジュ乗数の初期推定解
            Xk = Xk1

        # 終端コストのパラメータ 
        x_ref = [0.0, math.pi, 0.0, 0.0]  # 目標状態
        Q = [1.0, 10.0, dy_weight, dth_weight] # 状態への重み (学習パラメータ)
        Vf = 0                            # 終端コスト
        for i in range(nx):
            Vf += 0.5 * Q[i] * (Xk[i]-x_ref[i])**2 
        for i in range(nu):
            Vf += 0.5 * R[i] * Uk[i]**2
        J = J + Vf                        # コスト関数に終端コストを追加

        self.J = J
        self.w = vertcat(*w)
        self.g = vertcat(*g)
        self.x = w0
        self.lam_x = lam_x0
        self.lam_g = lam_g0
        self.lbx = lbw
        self.ubx = ubw
        self.lbg = lbg
        self.ubg = ubg

        # 非線形計画問題(NLP)
        self.nlp = {'f': self.J, 'x': self.w, 'g': self.g} 
        # Ipopt ソルバー，最小バリアパラメータを0.1，最大反復回数を5, ウォームスタートをONに
        self.solver = nlpsol('solver', 'ipopt', self.nlp, {'calc_lam_p':True, 'calc_lam_x':True, 'print_time':False, 'ipopt':{'max_iter':max_ipopt_iter, 'mu_min':0.1, 'warm_start_init_point':'yes', 'print_level':0, 'print_timing_statistics':'no'}}) 


    def init(self, x0=None):
        if x0 is not None:
            # 初期状態についての制約を設定
            self.lbx[0:4] = x0
            self.ubx[0:4] = x0
        # primal variables (x) と dual variables（ラグランジュ乗数）の初期推定解も与えつつ solve（warm start）
        sol = self.solver(x0=self.x, lbx=self.lbx, ubx=self.ubx, lbg=self.lbg, ubg=self.ubg, lam_x0=self.lam_x, lam_g0=self.lam_g)
        # 次の warm start のために解を保存
        self.x = sol['x'].full().flatten()
        self.lam_x = sol['lam_x'].full().flatten()
        self.lam_g = sol['lam_g'].full().flatten()

    def solve(self, x0):
        # 初期状態についての制約を設定
        nx = x0.shape[0]
        self.lbx[0:nx] = x0
        self.ubx[0:nx] = x0
        # primal variables (x) と dual variables（ラグランジュ乗数）の初期推定解も与えつつ solve（warm start）
        sol = self.solver(x0=self.x, lbx=self.lbx, ubx=self.ubx, lbg=self.lbg, ubg=self.ubg, lam_x0=self.lam_x, lam_g0=self.lam_g)
        # 次の warm start のために解を保存
        self.x = sol['x'].full().flatten()
        self.lam_x = sol['lam_x'].full().flatten()
        self.lam_g = sol['lam_g'].full().flatten()
        return self.x[4] # 制御入力を return


# Closed-loop シミュレーション
def simulation(train_params):
    sim_time = 10.0 # 10秒間のシミュレーション
    sampling_time = 0.0025 # 0.0025秒（2.5ms）のサンプリング周期
    sim_steps = math.floor(sim_time/sampling_time)
    xs = []
    us = []
    cartpole = CartPole()
    mpc = MPC(train_params)
    mpc.init()
    x = np.zeros(4)
    for step in range(sim_steps):
        if step%(1/sampling_time)==0:
            print('t =', step*sampling_time)
        u = mpc.solve(x)
        xs.append(x)
        us.append(u)
        x1 = x + sampling_time * np.array(cartpole.dynamics(x, u))
        x = x1
    return xs, us


# trial (シミュレーション) の評価
def objective(trial):
    T = trial.suggest_uniform('T', 0.1, 5.0)
    N = trial.suggest_int('N', 1, 50)
    dy_weight = trial.suggest_loguniform('dy_weight', 1.0e-08, 1.0)
    dth_weight = trial.suggest_loguniform('dth_weight', 1.0e-08, 1.0)
    u_weight = trial.suggest_loguniform('u_weight', 1.0e-08, 1.0)
    max_ipopt_iter = trial.suggest_int('max_ipopt_iter', 1, 20)
    train_params = TrainParams(T, N, dy_weight, dth_weight, u_weight, max_ipopt_iter)
    xs, us = simulation(train_params)
    # シミュレーション結果からトライアルのスコアを計算
    score = 0
    x_ref = [0.0, math.pi, 0.0, 0.0]  # 目標状態
    Q = [1.0, 10.0, 0, 0] # 真の（？）（正則化のない）評価関数
    for x in xs:
        for i in range(4):
            score += 0.5 * Q[i] * (x[i]-x_ref[i])**2 
    # 計算時間についてもスコアに加算
    # 本当はIpoptの計算時間を直接評価したいがやり方が分からないため，
    # 最大反復回数とホライゾン分割数からスコアを評価
    score += 100.0 * max_ipopt_iter * N
    print('score: %1.3f' %score)
    return score


# 学習（トライアル100回）
study = optuna.create_study(study_name='mpc_param_tune', 
                            storage='sqlite:///./mpc_param_tune.db',
                            load_if_exists=True)
study.optimize(objective, n_trials=100)

# 学習結果
print('Best parameters: ', study.best_params)


# 最適パラメータでシミュレーション
best_params = TrainParams()
best_params.T = study.best_params['T']
best_params.N = study.best_params['N']
best_params.dy_weight = study.best_params['dy_weight']
best_params.dth_weight = study.best_params['dth_weight'] 
best_params.u_weight = study.best_params['u_weight']
best_params.max_ipopt_iter = study.best_params['max_ipopt_iter']
xs, us = simulation(best_params)

# シミュレーション結果をプロット
sim_time = 10.0 # 10秒間のシミュレーション
sampling_time = 0.0025 # 0.0025秒（2.5ms）のサンプリング周期
sim_steps = math.floor(sim_time/sampling_time)
xs1 = [x[0] for x in xs]
xs2 = [x[1] for x in xs]
xs3 = [x[2] for x in xs]
xs4 = [x[3] for x in xs]
tgrid = [sampling_time*k for k in range(sim_steps)]

import matplotlib.pyplot as plt
plt.figure(1)
plt.clf()
plt.plot(tgrid, xs1, '--')
plt.plot(tgrid, xs2, '-')
plt.plot(tgrid, xs3, '-')
plt.plot(tgrid, xs4, '-')
plt.step(tgrid, us, '-.')
plt.xlabel('t')
plt.legend(['y(x1)','th(x2)', 'dy(x3)', 'dth(x4)','u'])
plt.grid()
plt.show()