SuccessiveConvexificationFreeFinalTime/main.py at master · EmbersArc/SuccessiveConvexificationFreeFinalTime · 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
from time import time
import numpy as np

from parameters import *
from discretization import Integrator
from scproblem import SCProblem
from utils import format_line, save_arrays

from models.model_6dof import Model_6DoF
from models.model_6dof_plot import plot3d

"""
Python implementation of 'Successive Convexification for 6-DoF Mars Rocket Powered Landing with Free-Final-Time' paper
by Michael Szmuk and Behçet Açıkmeşe.

Implementation by Sven Niederberger (s-niederberger@outlook.com)
"""

m = Model_6DoF()

# state and input
X = np.empty(shape=[m.n_x, K])
U = np.empty(shape=[m.n_u, K])

# INITIALIZATION--------------------------------------------------------------------------------------------------------
sigma = m.t_f_guess
X, U = m.initialize_trajectory(X, U)

# START SUCCESSIVE CONVEXIFICATION--------------------------------------------------------------------------------------
all_X = [X]
all_U = [U]

integrator = Integrator(m, K)
problem = SCProblem(m, K)

last_linear_cost = None

converged = False
for it in range(iterations):
    t0_it = time()
    print('-' * 50)
    print('-' * 18 + f' Iteration {str(it + 1).zfill(2)} ' + '-' * 18)
    print('-' * 50)

    t0_tm = time()
    A_bar, B_bar, C_bar, S_bar, z_bar = integrator.calculate_discretization(X, U, sigma)
    print(format_line('Time for transition matrices', time() - t0_tm, 's'))

    problem.set_parameters(A_bar=A_bar, B_bar=B_bar, C_bar=C_bar, S_bar=S_bar, z_bar=z_bar,
                           X_last=X, U_last=U, sigma_last=sigma,
                           weight_sigma=w_sigma, weight_nu=w_nu, weight_delta=w_delta, weight_delta_sigma=w_delta_sigma)

    while True:
        info = problem.solve(verbose=verbose_solver, solver=solver)
        print(format_line('Solver Error', info['solver_error']))

        X = problem.get_variable('X')
        U = problem.get_variable('U')
        sigma = problem.get_variable('sigma')

        delta_norm = problem.get_variable('delta_norm')
        sigma_norm = problem.get_variable('sigma_norm')
        nu_norm = np.linalg.norm(problem.get_variable('nu'), np.inf)

        print(format_line('delta_norm', delta_norm))
        print(format_line('sigma_norm', sigma_norm))
        print(format_line('nu_norm', nu_norm))

        print(format_line('flight time', sigma))

        if delta_norm < 1e-3 and sigma_norm < 1e-3 and nu_norm < 1e-7:
            converged = True

        w_delta *= 1.5

        problem.set_parameters(weight_delta=w_delta)

        break

    print('')
    print(format_line('Time for iteration', time() - t0_it, 's'))
    print('')

    all_X.append(X)
    all_U.append(U)

    if converged:
        print(f'Converged after {it + 1} iterations.')
        break

all_X = np.stack(all_X)
all_U = np.stack(all_U)

# save trajectory to file for visualization
save_arrays('output/trajectory/', {'X': all_X, 'U': all_U, 'sigma': sigma})

# plot trajectory
plot3d(all_X, all_U)