bilevel_design_opt/compute_bilevel_design_opt_performance.py at main · song970407/bilevel_design_opt · 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
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
import os
import argparse
import pickle

import numpy as np
import torch
import yaml

from src.env.HeatDiffusion import HeatDiffusionSystem
from src.solver.optimal_control import OptimalControl
from src.utils.get_target import generate_target_trajectory
from src.utils.fix_seed import fix_seed

fix_seed()
DEVICE = 'cuda:0' if torch.cuda.is_available() else 'cpu'
if not os.path.exists('bilevel_opt_result/optimal'):
    os.makedirs('bilevel_opt_result/optimal')


def run_optimal_control(mpc_config, env_config, data_generation_config, data_preprocessing_config, state_pos, action_pos):
    ridge_coefficient = mpc_config['ridge_coefficient']
    smoothness_coefficient = mpc_config['smoothness_coefficient']
    target_values_list = mpc_config['target_values_list']
    target_times_list = mpc_config['target_times_list']
    max_iter = mpc_config['max_iter']
    loss_threshold = mpc_config['loss_threshold']
    opt_config = mpc_config['opt_config']
    scheduler_config = mpc_config['scheduler_config']

    num_cells = env_config['num_cells']
    dt = env_config['dt']
    epsilon = env_config['epsilon']
    domain_range = env_config['domain_range']
    action_min = data_generation_config['action_bound'][0]
    action_max = data_generation_config['action_bound'][1]
    state_scaler = data_preprocessing_config['state_scaler']
    action_scaler = data_preprocessing_config['action_scaler']

    num_states = state_pos.shape[0]
    num_actions = action_pos.shape[0]

    u_min = action_min
    u_max = action_max
    is_logging = True

    target_list = []
    for (target_values, target_times) in zip(target_values_list, target_times_list):
        target = np.array(generate_target_trajectory(target_values, target_times))
        target = np.reshape(target, newshape=(-1, 1))
        target = np.concatenate([target for _ in range(num_states)], axis=1)
        target = torch.from_numpy(target).float().to(DEVICE)
        target_list.append(target)
    receding_horizon = target_list[0].shape[0]

    env = HeatDiffusionSystem(num_cells=num_cells,
                              dt=dt,
                              epsilon=epsilon,
                              domain_range=domain_range,
                              action_min=action_min,
                              action_max=action_max,
                              state_pos=state_pos,
                              action_pos=action_pos)
    env.reset()

    optimal_controller = OptimalControl(env, num_states, num_actions, receding_horizon, ridge_coefficient,
                                        smoothness_coefficient, u_min, u_max, max_iter, loss_threshold, is_logging,
                                        DEVICE, opt_config, scheduler_config)
    trajectory_x = []
    trajectory_u = []
    trajectory_log = []

    for (i, target) in enumerate(target_list):
        print('Now target number {}'.format(i))
        optimal_us, log = optimal_controller.solve(target)
        trajectory_log.append(log)
        x_traj = []
        x_traj.append(np.zeros((num_states)))
        u_traj = []
        env.reset()
        for optimal_u in optimal_us:
            x_traj.append(env.step(optimal_u))
            u_traj.append(optimal_u)
        x_traj = np.stack(x_traj)
        u_traj = np.stack(u_traj)
        trajectory_x.append(x_traj)
        trajectory_u.append(u_traj)
    return trajectory_x, trajectory_u, trajectory_log


if __name__ == '__main__':
    env_config = yaml.safe_load(open('config/env/env_config.yaml', 'r'))
    data_generation_config = yaml.safe_load(open('config/data/data_generation_config.yaml', 'r'))
    data_preprocessing_config = yaml.safe_load(open('config/data/data_preprocessing_config.yaml', 'r'))
    target = pickle.load(open('data/bilevel_design_opt/target.pkl', 'rb'))

    mpc_config = {
        'ridge_coefficient': 0,
        'smoothness_coefficient': 0,
        'target_values_list': target['target_values'],
        'target_times_list': target['target_times'],
        'max_iter': 200,
        'loss_threshold': 1e-9,
        'opt_config': {'lr': 2e-0},
        'scheduler_config': {'patience': 5, 'factor': 0.5, 'min_lr': 1e-4}
    }

    parser = argparse.ArgumentParser()
    # parser.add_argument('--num_x', type=int, default=3)
    # parser.add_argument('--num_heaters', type=int, default=5)
    parser.add_argument('--solver_name', type=str, default='implicit')
    parser.add_argument('--model_name', type=str, default='Linear')
    args = parser.parse_args()

    solver_name = args.solver_name
    model_name = args.model_name

    num_x_list = [3, 4, 5]
    num_heaters_list = [5, 10, 15, 20]

    if not os.path.exists('bilevel_opt_result/optimal/{}_{}'.format(solver_name, model_name)):
        os.makedirs('bilevel_opt_result/optimal/{}_{}'.format(solver_name, model_name))

    for num_x in num_x_list:
        for num_heaters in num_heaters_list:
            opt_result = pickle.load(open('bilevel_opt_result/{}_{}/{}_{}.pkl'.format(solver_name, model_name, num_x, num_heaters), 'rb'))
            state_pos = pickle.load(open('data/bilevel_design_opt/problem_{}_{}.pkl'.format(num_x, num_heaters), 'rb'))['state_pos'][0]
            num_repeats = opt_result['opt_log']['total_loss_trajectory'].shape[1]
            for i in range(num_repeats):
                print('Now {}, {}, {}'.format(num_x, num_heaters, i))
                if os.path.isfile('bilevel_opt_result/optimal/{}_{}/{}_{}_{}.pkl'.format(solver_name, model_name, num_x, num_heaters, i)):
                    continue
                best_idx = np.argmin(opt_result['opt_log']['total_loss_trajectory'][:, i])
                action_pos = opt_result['opt_log']['position_trajectory'][best_idx, i]
                x_trajectory_list, u_trajectory_list, log_trajectory_list = run_optimal_control(mpc_config,
                                                                                                env_config,
                                                                                                data_generation_config,
                                                                                                data_preprocessing_config,
                                                                                                state_pos,
                                                                                                action_pos)
                # pickle.dump(mpc_config, open('bilevel_opt_result/optimal/{}/mpc_config_{}_{}.pkl'.format(solver_name, num_x, num_heaters), 'wb'))
                optimal_result = {
                    'x_trajectory_list': x_trajectory_list,
                    'u_trajectory_list': u_trajectory_list,
                    'log_trajectory_list': log_trajectory_list
                }
                pickle.dump(optimal_result, open('bilevel_opt_result/optimal/{}_{}/{}_{}_{}.pkl'.format(solver_name, model_name, num_x, num_heaters, i), 'wb'))