コード例 #1
0
ファイル: example.py プロジェクト: ADVRHumanoids/ilqr
def doit():

    # state vector is (x, y, theta)
    x = cs.SX.sym('x', 3)

    # control vector is (v, omega)
    u = cs.SX.sym('u', 2)

    # explicit variables for readability
    v = u[0]
    omega = u[1]
    theta = x[2]

    # unicycle kinematic model
    xdot = cs.vertcat(v*cs.cos(theta),
                      v*cs.sin(theta),
                      omega)

    # intermediate cost, final cost, final constraint
    l = cs.sumsqr(u)
    xf_des = np.array([0, 1, 0])
    lf = 1000.0 * cs.sumsqr(x - xf_des)
    hf = x - xf_des

    # create solver
    solver = ilqr.IterativeLQR(x=x,
                               u=u,
                               xdot=xdot,
                               dt=0.1,
                               N=100,
                               intermediate_cost=l,
                               final_cost=lf,
                               final_constraint=None)
    # solve
    solver.randomizeInitialGuess()
    solver.solve(50)

    plt.figure()
    plt.title('Total cost')
    plt.semilogy(solver._dcost, '--s')
    plt.grid()

    plt.figure()
    plt.title('State trajectory')
    plt.plot(solver._state_trj)
    plt.legend(['x', 'y', 'theta'])
    plt.grid()

    plt.figure()
    plt.title('Control trajectory')
    plt.plot(solver._ctrl_trj)
    plt.legend(['v', 'omega'])
    plt.grid()

    plt.show()
コード例 #2
0
ファイル: BaseOpti.py プロジェクト: BDEvan5/TestAlgorithms
def BaseOptiDX():
    pathpoints = 30
    ref_path = {}
    ref_path['x'] = 5 * np.sin(np.linspace(0, 2 * np.pi, pathpoints + 1))
    ref_path['y'] = np.linspace(1, 2, pathpoints + 1)**2 * 10
    wp = ca.horzcat(ref_path['x'], ref_path['y']).T

    N = 5
    N1 = N - 1

    wpts = np.array(wp[:, 0:N])

    x = ca.MX.sym('x', N)
    dx = ca.MX.sym('dx', N - 1)

    nlp = {\
        'x': ca.vertcat(x, dx),
        'f': ca.sumsqr(x - wpts[0, :]) + ca.sumsqr(dx),
        'g': ca.vertcat(\
            x[1:] - (x[:-1] + dx),

            x[0] - wpts[0, 0],
            )\
        }

    S = ca.nlpsol('S', 'ipopt', nlp)

    # x0
    x00 = wpts[0, :]
    dx00 = (x00[1:] - x00[:-1])

    x0 = ca.vertcat(x00, dx00)

    lbx = [0] * N + [0] * N1
    ubx = [ca.inf] * N + [10] * N1

    lbg = [0] * N
    ubg = [0] * N

    # solve
    r = S(x0=x0, lbx=lbx, ubx=ubx, lbg=lbg, ubg=ubg)
    x_opt = r['x']

    x = np.array(x_opt[:N])
    xds = np.array(x_opt[N:N + N1])

    print(f"xs: {x.T}")
    print(f"dx's: {xds.T}")

    plt.figure(1)
    plt.plot(wpts[0, :], np.ones_like(wpts[0, :]), 'o', markersize=12)

    plt.plot(x, np.ones_like(x), '+', markersize=20)

    plt.show()
コード例 #3
0
def smooth_track(track):
    N = len(track)
    xs = track[:, 0]
    ys = track[:, 1]

    th1_f = ca.MX.sym('y1_f', N-2)
    th2_f = ca.MX.sym('y2_f', N-2)
    x_f = ca.MX.sym('x_f', N)
    y_f = ca.MX.sym('y_f', N)

    real = ca.Function('real', [th1_f, th2_f], [ca.cos(th1_f)*ca.cos(th2_f) + ca.sin(th1_f)*ca.sin(th2_f)])
    im = ca.Function('im', [th1_f, th2_f], [-ca.cos(th1_f)*ca.sin(th2_f) + ca.sin(th1_f)*ca.cos(th2_f)])

    sub_cmplx = ca.Function('a_cpx', [th1_f, th2_f], [ca.atan(im(th1_f, th2_f)/real(th1_f, th2_f))])

    d_th = ca.Function('d_th', [x_f, y_f], [ca.if_else(ca.fabs(x_f[1:] - x_f[:-1]) < 0.01 ,ca.atan((y_f[1:] - y_f[:-1])/(x_f[1:] - x_f[:-1])), 10000)])

    x = ca.MX.sym('x', N)
    y = ca.MX.sym('y', N)
    th = ca.MX.sym('th', N-1)

    B = 5
    nlp = {\
        'x': ca.vertcat(x, y, th),
        'f': ca.sumsqr(sub_cmplx(th[1:], th[:-1])) + B* (ca.sumsqr(x-xs) + ca.sumsqr(y-ys)),
        # 'f':  B* (ca.sumsqr(x-xs) + ca.sumsqr(y-ys)),
        'g': ca.vertcat(\
                th - d_th(x, y),
                x[0] - xs[0], y[0]- ys[0],
                x[-1] - xs[-1], y[-1]- ys[-1],
            )\
        }

    S = ca.nlpsol('S', 'ipopt', nlp)
    # th0 = [lib.get_bearing(track[i, 0:2], track[i+1, 0:2]) for i in range(N-1)]
    th0 = d_th(xs, ys)

    x0 = ca.vertcat(xs, ys, th0)

    lbx = [0] *2* N + [-np.pi]*(N -1)
    ubx = [100] *2 * N + [np.pi]*(N-1) 

    r = S(x0=x0, lbg=0, ubg=0, lbx=lbx, ubx=ubx)

    print(f"Solution found")
    x_opt = r['x']

    xs_new = np.array(x_opt[:N])
    ys_new = np.array(x_opt[N:2*N])

    track[:, 0] = xs_new[:, 0]
    track[:, 1] = ys_new[:, 0]

    return track
コード例 #4
0
    def init_ocp(self):
        ocp = Ocp(T=FreeTime(20.0))

        self.x = ocp.state()
        self.y = ocp.state()
        self.theta = ocp.state()

        self.v = ocp.control()
        self.theta_dot = ocp.control()

        ocp.set_der(self.x, sin(self.theta) * self.v)
        ocp.set_der(self.y, cos(self.theta) * self.v)
        ocp.set_der(self.theta, self.theta_dot)

        self.X_0 = ocp.parameter(self.nx)

        X = vertcat(self.x, self.y, self.theta)
        ocp.subject_to(ocp.at_t0(X) == self.X_0)

        max_v = 5
        ocp.subject_to(1 <= (self.v <= max_v))
        th_dot_lim = 0.5
        ocp.subject_to(-th_dot_lim <= (self.theta_dot <= th_dot_lim))
        #ocp.subject_to( -0.3 <= (ocp.der(V) <= 0.3) )

        # Minimal time
        # ocp.add_objective(0.50*ocp.T)

        self.waypoints = ocp.parameter(2, grid='control')
        self.waypoint_last = ocp.parameter(2)
        p = vertcat(self.x, self.y)

        ocp.add_objective(ocp.sum(sumsqr(p - self.waypoints), grid='control'))
        ocp.add_objective(sumsqr(ocp.at_tf(p) - self.waypoint_last))

        options = {"ipopt": {"print_level": 5}}
        options["expand"] = True
        options["print_time"] = False
        ocp.solver('ipopt', options)

        ocp.method(
            MultipleShooting(N=self.N,
                             M=1,
                             intg='rk',
                             grid=FreeGrid(min=0.05, max=2)))

        ocp.set_initial(self.x, 0)
        ocp.set_initial(self.y, 0)
        ocp.set_initial(self.theta, 0)
        ocp.set_initial(self.theta_dot, 0)
        ocp.set_initial(self.v, 2)

        self.ocp = ocp
コード例 #5
0
def generate_brent(n=2, a=10., b=10., func_opts={}, data_type=cs.SX):
    if n != 2:
        raise ValueError("brent is only defined for n=2")
    x = data_type.sym("x", n)
    f = (x[0] + a)**2 + (x[1] + b)**2 + cs.exp(-cs.sumsqr(x))
    func = cs.Function("brent", [x], [f], ["x"], ["f"], func_opts)
    return func, [[-20., 0.]] * n, [[-10.] * n]
コード例 #6
0
def robust_control_stages(ocp, delta, t0):
    """
  Returns
  -------
  stage: :obj: `~rockit.stage.Stage
  x : :obj: `~casadi.MX

  """
    stage = ocp.stage(t0=t0, T=1)
    x = stage.state(2)
    u = stage.state()
    der_state = vertcat(x[1],
                        -0.1 * (1 - x[0]**2 + delta) * x[1] - x[0] + u + delta)
    stage.set_der(x, der_state)
    stage.set_der(u, 0)
    stage.subject_to(u <= 40)
    stage.subject_to(u >= -40)
    stage.subject_to(x[0] >= -0.25)
    L = stage.variable()
    stage.add_objective(L)
    stage.subject_to(L >= sumsqr(x[0] - 3))
    bound = lambda t: 2 + 0.1 * cos(10 * t)
    stage.subject_to(x[0] <= bound(stage.t))
    stage.method(MultipleShooting(N=20, M=1, intg='rk'))

    return stage, x, stage.at_t0(u)
コード例 #7
0
def generate_egg_crate(n=2, a=25., func_opts={}, data_type=cs.SX):
    if n != 2:
        raise ValueError("egg_crate is only defined for n=2")
    x = data_type.sym("x", n)
    f = cs.sumsqr(x) + a * (cs.sin(x[0])**2 + cs.sin(x[1])**2)
    func = cs.Function("egg_crate", [x], [f], ["x"], ["f"], func_opts)
    return func, [[-5., 5.]] * n, [[0., 0.]]
コード例 #8
0
def generate_griewank(n=2, a=1., b=4000., func_opts={}, data_type=cs.SX):
    x = data_type.sym("x", n)
    product_part = 1.
    for i in range(n):
        product_part = product_part * cs.cos(x[i] / cs.sqrt(i + 1))
    f = a + (1. / b) * cs.sumsqr(x) - product_part
    func = cs.Function("griewank", [x], [f], ["x"], ["f"], func_opts)
    return func, [[-600., 600.]] * n, [[0.] * n]
コード例 #9
0
def generate_deckkers_aarts(n=2, a=1e5, b=1e-5, func_opts={}, data_type=cs.SX):
    if n != 2:
        raise ValueError("deckkers_aarts is only defined for n=2")
    x = data_type.sym("x", n)
    sqr = cs.sumsqr(x)
    f = a * x[0]**2 + x[1]**2 - sqr**2 + b * sqr**4
    func = cs.Function("deckkers_aarts", [x], [f], ["x"], ["f"], func_opts)
    minima = [[0., 15.], [0., -15.]]
    return func, [[-20., 20.], [-20., 20.]], minima
コード例 #10
0
def generate_bartelsconn(n=2, func_opts={}, data_type=cs.SX):
    if n != 2:
        raise ValueError("bartelsconn is only defined for n=2")
    x = data_type.sym("x", n)
    f = cs.norm_1(cs.sumsqr(x) + x[0] * x[1])
    f += cs.norm_1(cs.sin(x[0]))
    f += cs.norm_1(cs.cos(x[1]))
    func = cs.Function("bartelsconn", [x], [f], ["x"], ["f"], func_opts)
    return func, [[-500., 500.]] * n, [[0.] * n]
コード例 #11
0
    def create_objective(self, model):
        self.obj_1 = sum(
            w / len(ov) * ca.sumsqr(self.densities * (ov - cm @ ca.interp1d(
                model.observation_times,
                om(model.observation_times, model.ps, *
                   (model.xs[j] for j in oj)), self.observation_times)))
            for om, oj, ov, w, cm in zip(
                self.observation_model, self.observation_vector,
                self.observations, self.weightings, self.collocation_matrices))
        self.obj_2 = sum(
            ca.sumsqr((model.xdash[:, i] - model.model(
                model.observation_times, *model.cs, *model.ps)[:, i]))
            for i in range(model.s)) / model.n

        self.regularisation = ca.sumsqr(
            (ca.vcat(model.ps) - ca.vcat(self.regularisation_vector)) /
            (1 + ca.vcat(self.regularisation_vector)))

        self.objective = self.obj_1 + self.rho * self.obj_2 + self.alpha * self.regularisation
コード例 #12
0
ファイル: RockitClass.py プロジェクト: BDEvan5/TestAlgorithms
    def init_mpc(self):
        # Specify ODE
        self.ocp.set_der(self.x, self.V * cos(self.theta))
        self.ocp.set_der(self.y, self.V * sin(self.theta))
        self.ocp.set_der(self.theta, self.V / L * tan(self.delta))

        # Initial constraints
        self.ocp.subject_to(self.ocp.at_t0(self.X) == self.X_0)

        # Initial guess
        self.ocp.set_initial(self.x, 0)
        self.ocp.set_initial(self.y, 0)
        self.ocp.set_initial(self.theta, 0)

        self.ocp.set_initial(self.V, 0.5)

        # Path constraints
        max_v = 5
        self.ocp.subject_to(0 <= (self.V <= max_v))
        #ocp.subject_to( -0.3 <= (ocp.der(V) <= 0.3) )
        self.ocp.subject_to(-pi / 6 <= (self.delta <= pi / 6))

        # Minimal time
        self.ocp.add_objective(0.50 * self.ocp.T)

        self.ocp.add_objective(
            self.ocp.sum(sumsqr(self.p - self.waypoints), grid='control'))
        self.ocp.add_objective(
            sumsqr(self.ocp.at_tf(self.p) - self.waypoint_last))

        # Pick a solution method
        options = {"ipopt": {"print_level": 0}}
        options["expand"] = True
        options["print_time"] = False
        self.ocp.solver('ipopt', options)

        # Make it concrete for this ocp
        self.ocp.method(
            MultipleShooting(N=N,
                             M=1,
                             intg='rk',
                             grid=FreeGrid(min=0.05, max=2)))
コード例 #13
0
def generate_happy_cat(n=2,
                       a=0.125,
                       b=0.5,
                       c=0.5,
                       func_opts={},
                       data_type=cs.SX):
    x = data_type.sym("x", n)
    nx2 = cs.dot(x, x)
    f = ((nx2 - n)**2)**a + (1. / n) * (b * nx2 + cs.sumsqr(x)) + c
    func = cs.Function("happy_cat", [x], [f], ["x"], ["f"], func_opts)
    return func, [[-2., 2.]] * n, [[-1] * n]
コード例 #14
0
def generate_salomon(n=2,
                     a=1.,
                     b=2 * cs.np.pi,
                     c=0.1,
                     func_opts={},
                     data_type=cs.SX):
    x = data_type.sym("x", n)
    sumsq = cs.sumsqr(x)
    f = a - cs.cos(b * cs.sqrt(sumsq)) + c * cs.sqrt(sumsq)
    func = cs.Function("salomon", [x], [f], ["x"], ["f"], func_opts)
    return func, [[-10., 10.]] * n, [[0.] * n]
コード例 #15
0
def optCtrl(optQueue,opt2Queue,P,MU,S,H,tgt,optdt):
    sold=None
    
    opt=qnTransient()
    opt.buildOpt(MU, P, S,optdt,tgt,H)
    
    while(True):
            
        q=optQueue.get()
        opt.model.set_value(opt.initX,q)
        
        obj=0
        for h in range(H):
            obj+=(opt.stateVar[1,h]-tgt/MU[1,0]*opt.tVar[1,h])**2
            #obj+=opt.E_abs[0,h]
        if(sold is not None):
            opt.model.minimize(obj+0.0*casadi.sumsqr(opt.sVar-sold)+0.000*casadi.sumsqr(opt.sVar))
        else:
            opt.model.minimize(obj)
        
        sol=opt.model.solve()
        opt2Queue.put(sol.value(opt.sVar))
コード例 #16
0
ファイル: falcon9_rtls.py プロジェクト: janismac/yamtof
def add_target_orbit_constraint(mocp, p, phase_insertion):
    start = lambda a: mocp.start(a)
    end = lambda a: mocp.end(a)

    # Target orbit - soft constraints
    # 3 parameter orbit: SMA,ECC,INC
    #
    # h_target_magnitude - h_final_magnitude == 0
    # h_target_z - h_final_z == 0
    # c_target - c_final == 0

    r_final_ECEF = casadi.vertcat(end(phase_insertion.rx),
                                  end(phase_insertion.ry),
                                  end(phase_insertion.rz))
    v_final_ECEF = casadi.vertcat(end(phase_insertion.vx),
                                  end(phase_insertion.vy),
                                  end(phase_insertion.vz))

    # Convert to ECI (z rotation is omitted, because LAN is free)
    r_f = r_final_ECEF
    v_f = (
        v_final_ECEF +
        casadi.cross(casadi.SX([0, 0, p.body_rotation_speed]), r_final_ECEF))

    h_final = casadi.cross(r_f, v_f)
    h_final_mag = casadi.norm_2(h_final)
    h_final_z = h_final[2]
    c_final = casadi.sumsqr(v_f) / 2 - 1.0 / casadi.norm_2(r_f)

    defect_h_magnitude = h_final_mag - p.target_orbit_h_magnitude
    defect_h_z = h_final_z - p.target_orbit_h_z
    defect_c = c_final - p.target_orbit_c

    slack_h_magnitude = mocp.add_variable('slack_h_magnitude', init=3.0)
    slack_h_z = mocp.add_variable('slack_h_z', init=3.0)
    slack_c = mocp.add_variable('slack_c', init=3.0)

    mocp.add_constraint(slack_h_magnitude > 0)
    mocp.add_constraint(slack_h_z > 0)
    mocp.add_constraint(slack_c > 0)

    mocp.add_constraint(defect_h_magnitude < slack_h_magnitude)
    mocp.add_constraint(defect_h_magnitude > -slack_h_magnitude)
    mocp.add_constraint(defect_h_z < slack_h_z)
    mocp.add_constraint(defect_h_z > -slack_h_z)
    mocp.add_constraint(defect_c < slack_c)
    mocp.add_constraint(defect_c > -slack_c)

    mocp.add_objective(100.0 * slack_h_magnitude)
    mocp.add_objective(100.0 * slack_h_z)
    mocp.add_objective(100.0 * slack_c)
コード例 #17
0
    def solve(self, x0):
        xs = ca.vcat([state for state in self.states])
        us = ca.vcat([control for control in self.controls])

        dyns = ca.vcat([
            var[1:] - (var[:-1] + self.state_der[var] * self.dt)
            for var in self.states
        ])
        cons = ca.vcat(
            [state[0] - x0[i] for i, state in enumerate(self.states)])

        obs = ca.vcat([(o - self.objectives[o]) * self.o_scales[o]
                       for o in self.objectives.keys()])

        nlp = {\
            'x': ca.vertcat(xs, us, self.dt),
            'f': ca.sumsqr(obs),
            'g': ca.vertcat(dyns, cons)
            }

        n_g = nlp['g'].shape[0]
        self.lbg = [0] * n_g
        self.ubg = [0] * n_g

        x00 = ca.vcat([self.initial[state] for state in self.states])
        u00 = ca.vcat([self.initial[control] for control in self.controls])
        x0 = ca.vertcat(x00, u00, self.initial[self.dt])

        # S = ca.nlpsol('S', 'ipopt', nlp, {'ipopt':{'print_level':0}})
        S = ca.nlpsol('S', 'ipopt', nlp, {'ipopt': {'print_level': 5}})
        r = S(x0=x0, lbx=self.lbx, ubx=self.ubx, lbg=self.lbg, ubg=self.ubg)
        x_opt = r['x']

        n_state_vars = len(self.states) * self.n
        n_control_vars = len(self.controls) * self.n

        states = np.array(x_opt[:n_state_vars])
        controls = np.array(x_opt[n_state_vars:n_state_vars + n_control_vars])
        times = np.array(x_opt[-self.n + 1:])

        for i, state in enumerate(self.states):
            self.set_inital(state, states[i * self.n:self.n * (i + 1)])

        for i, control in enumerate(self.controls):
            self.set_inital(control,
                            controls[(self.n - 1) * i:(i + 1) * (self.n - 1)])

        self.set_inital(self.dt, times)

        return states, controls, times
コード例 #18
0
def unit_quat(q):
    """
    Normalizes a quaternion to be unit modulus.
    :param q: 4-dimensional numpy array or CasADi object
    :return: the unit quaternion in the same data format as the original one
    """

    if isinstance(q, np.ndarray):
        # if (q == np.zeros(4)).all():
        #     q = np.array([1, 0, 0, 0])
        q_norm = np.sqrt(np.sum(q**2))
    else:
        q_norm = cs.sqrt(cs.sumsqr(q))
    return 1 / q_norm * q
コード例 #19
0
def generate_drop_wave(n=2,
                       a=1.,
                       b=12.,
                       c=0.5,
                       d=2.,
                       func_opts={},
                       data_type=cs.SX):
    if n != 2:
        raise ValueError("drop_wave is only defined for n=2")
    x = data_type.sym("x", n)
    sqr = cs.sumsqr(x)
    f = -(a + cs.cos(b * cs.sqrt(sqr))) / (c * sqr + d)
    func = cs.Function("drop_wave", [x], [f], ["x"], ["f"], func_opts)
    return func, [[-5.2, 5.2]] * n, [[0., 0.]]
コード例 #20
0
def ModelTraining(model,
                  data,
                  initializations=10,
                  BFR=False,
                  p_opts=None,
                  s_opts=None):

    results = []

    for i in range(0, initializations):

        # in first run use initial model parameters (useful for online
        # training when only time for one initialization)
        if i > 0:
            model.Initialize()

        # Estimate Parameters on training data
        new_params = ModelParameterEstimation(model, data, p_opts, s_opts)

        # Assign new parameters to model
        model.Parameters = new_params

        # Evaluate on Validation data
        u_val = data['u_val']
        y_ref_val = data['y_val']
        init_state_val = data['init_state_val']

        # Loop over all experiments

        e = 0

        for j in range(0, u_val.shape[0]):

            # Simulate Model
            y = model.Simulation(init_state_val[j], u_val[j])
            y = np.array(y)

            e = e + cs.sumsqr(y_ref_val[j] - y)

        # Calculate mean error over all validation batches
        e = e / u_val.shape[0]
        e = np.array(e).reshape((1, ))

        # save parameters and performance in list
        results.append([e, new_params])

    results = pd.DataFrame(data=results, columns=['loss', 'params'])

    return results
コード例 #21
0
def gen(ocpVars):
    constrList = []
    M = ocpVars['u'].shape[1]

    for i in range(M):
        quat = ocpVars['u'][1:5, i]

        # norm has to be one
        g = ca.sumsqr(quat) - 1
        lbg = np.zeros(g.shape)
        ubg = np.zeros(g.shape)

        # create constraint tuple
        constrList += [(g, lbg, ubg)]

    return constrList
コード例 #22
0
ファイル: OCP.py プロジェクト: BDEvan5/TestAlgorithms
    def solve(self):
        xs = ca.vcat([state for state in self.states])
        us = ca.vcat([control for control in self.controls])

        dyns = ca.vcat([var[1:] - (var[:-1] + self.state_der[var] * self.dt) for var in self.states])
        cons = ca.vcat([cons[0] - self.constraints[cons] for cons in self.constraints.keys()])

        obs = ca.vcat([(o - self.objectives[o]) * self.o_scales[o] for o in self.objectives.keys()])

        nlp = {\
            'x': ca.vertcat(xs, us, self.dt),
            'f': ca.sumsqr(obs),
            'g': ca.vertcat(dyns, cons)
            }

        n_g = nlp['g'].shape[0]
        lbg = [0] * n_g
        ubg = [0] * n_g

        x00 = ca.vcat([self.initial[state] for state in self.states])
        u00 = ca.vcat([self.initial[control] for control in self.controls])
        x0 = ca.vertcat(x00, u00, self.initial[self.dt])

        x_mins = [self.min_lims[state] for state in self.states] * self.n
        x_maxs = [self.max_lims[state] for state in self.states] * self.n
        u_mins = [self.min_lims[control] for control in self.controls] * (self.n - 1)
        u_maxs = [self.max_lims[control] for control in self.controls] * (self.n - 1)

        lbx = x_mins + u_mins + list(np.zeros(self.n-1))
        ubx = x_maxs + u_maxs + list(np.ones(self.n-1) * self.max_t)
        
        S = ca.nlpsol('S', 'ipopt', nlp)
        r = S(x0=x0, lbx=lbx, ubx=ubx, lbg=lbg, ubg=ubg)
        x_opt = r['x']

        n_state_vars = len(self.states) * self.n
        n_control_vars = len(self.controls) * self.n

        states = np.array(x_opt[:n_state_vars])
        controls = np.array(x_opt[n_state_vars:n_state_vars + n_control_vars])
        times = np.array(x_opt[-self.n+1:])

        return states, controls, times
コード例 #23
0
ファイル: falcon9_rtls.py プロジェクト: janismac/yamtof
def solve_launch_direction_problem(p):
    # Find the approximate launch direction vector.
    # The launch direction vector is horizontal (orthogonal to the launch site position vector).
    # The launch direction vector points (approximately) in the direction of the orbit insertion velocity.
    mocp = MOCP()

    rx = mocp.add_variable('rx', init=p.launch_site_x)
    ry = mocp.add_variable('ry', init=p.launch_site_y)
    rz = mocp.add_variable('rz', init=p.launch_site_z)

    vx = mocp.add_variable('vx', init=0.01)
    vy = mocp.add_variable('vy', init=0.01)
    vz = mocp.add_variable('vz', init=0.01)

    mocp.add_objective((rx - p.launch_site_x)**2)
    mocp.add_objective((ry - p.launch_site_y)**2)
    mocp.add_objective((rz - p.launch_site_z)**2)

    r = casadi.vertcat(rx, ry, rz)
    v = casadi.vertcat(vx, vy, vz)
    h = casadi.cross(r, v)
    h_mag = casadi.norm_2(h)
    h_z = h[2]
    c = casadi.sumsqr(v) / 2 - 1.0 / casadi.norm_2(r)

    mocp.add_objective((h_mag - p.target_orbit_h_magnitude)**2)
    mocp.add_objective((h_z - p.target_orbit_h_z)**2)
    mocp.add_objective((c - p.target_orbit_c)**2)

    mocp.solve()
    v_soln = DM([mocp.get_value(vx), mocp.get_value(vy), mocp.get_value(vz)])
    launch_direction = normalize(v_soln)
    up_direction = normalize(
        DM([p.launch_site_x, p.launch_site_y, p.launch_site_z]))
    launch_direction = launch_direction - (
        launch_direction.T @ up_direction) * up_direction
    return launch_direction
コード例 #24
0
    def eval(self, x, w):
        '''
        evaluate the cost function, with casadi operation
        input:
        x: primal decision var
        w: RV, randomness, in ml: it's the data
        '''
        if len(x.shape)==1:
            x=x.reshape(-1,1)
        elif len(x.shape)==2:
            pass
        else:
            raise NotImplementedError

        if self.method =='boyd':
            '''
            Boyd & Vandenberghe book. Figure 6.15.

            min_x || A(w) - b0 ||^2 where
            A(w) := A0 + w * B0 and w is a scalar.
            '''
            # boy data set
            A0, B0, b0 = self.model  # model of the optimization problem
            if self.mode== 'casadi':
                from casadi import sumsqr
                cost_val = sumsqr((A0 + w * B0) @ x - b0)
            elif self.mode== 'cvxpy':
                import cvxpy as cp
                cost_val = cp.sum_squares((A0 + w * B0) @ x - b0)
            elif self.mode=='numpy':
                cost_val = np.sum(((A0 + w * B0) @ x - b0)**2)
            else:
                raise NotImplementedError
        else:
            raise NotImplementedError

        return cost_val
コード例 #25
0
			q_torso_dot_sol = var_dot_sol[16]

			print("Solver time is = " + str(hlp.time_taken))
			comp_time.append(hlp.time_taken)
			con_viol1 = hlp.optimal_slacks[1]
			con_viol2 = hlp.optimal_slacks[2]
			con_viol3 = hlp.optimal_slacks[3]
			con_viol4 = hlp.optimal_slacks[4]
			constraint_violations = cs.horzcat(constraint_violations, cs.vertcat(con_viol1, con_viol2, con_viol3, con_viol4))

		# print(con_viols)
		# print(q_opt)
		# print(s2_opt)
		number_non_zero = sum(cs.fabs(q_dot1_sol) >=  1e-3) + sum(cs.fabs(q_dot2_sol) >=  1e-3)
		q_zero_integral += number_non_zero*ts
		q_2_integral += (cs.sumsqr(q_dot1_sol) + cs.sumsqr(q_dot2_sol))*ts
		q_1_integral += sum(cs.fabs(q_dot1_sol).full() + cs.fabs(q_dot2_sol).full())*ts
		#compute q_1_integral
		# print(number_non_zero)
		#Update all the variables
		q_dot_opt = cs.vertcat(q_dot1_sol, q_dot2_sol, s_dot1_sol,  s_dot2_sol, q_torso_dot_sol)
		q_opt[0:17] += ts*q_dot_opt

		s1_opt = q_opt[14]
		if s1_opt >=1:
			print("Robot1 reached it's goal. Terminating")
			cool_off_counter += 1
			# break
			# if cool_off_counter >= 100: #ts*100 cooloff period for the robot to exactly reach it's goal
				# break
コード例 #26
0
    def compute_mprim(state_i, lattice, L, v, u_max, print_level):
        N = 75
        nx = 3
        Nc = 3
        mprim = []
        for state_f in lattice:
            x = ca.MX.sym('x', nx)
            u = ca.MX.sym('u')

            F = ca.Function('f', [x, u],
                            [v*ca.cos(x[2]), v*ca.sin(x[2]), v*ca.tan(u)/L])

            opti = ca.Opti()
            X = opti.variable(nx, N+1)
            pos_x = X[0, :]
            pos_y = X[1, :]
            ang_th = X[2, :]
            U = opti.variable(N, 1)
            T = opti.variable(1)
            dt = T/N

            opti.set_initial(T, 0.1)
            opti.set_initial(U, 0.0*np.ones(N))
            opti.set_initial(pos_x, np.linspace(state_i[0], state_f[0], N+1))
            opti.set_initial(pos_y, np.linspace(state_i[1], state_f[1], N+1))

            tau = ca.collocation_points(Nc, 'radau')
            C, _ = ca.collocation_interpolators(tau)

            Xc_vec = []
            for k in range(N):
                Xc = opti.variable(nx, Nc)
                Xc_vec.append(Xc)
                X_kc = ca.horzcat(X[:, k], Xc)
                for j in range(Nc):
                    fo = F(Xc[:, j], U[k])
                    opti.subject_to(X_kc@C[j+1] == dt*ca.vertcat(fo[0], fo[1], fo[2]))

                opti.subject_to(X_kc[:, Nc] == X[:, k+1]) 

            for k in range(N):
                opti.subject_to(U[k] <= u_max)
                opti.subject_to(-u_max <= U[k])

            opti.subject_to(T >= 0.001)
            opti.subject_to(X[:, 0] == state_i)
            opti.subject_to(X[:, -1] == state_f)

            alpha = 1e-2
            opti.minimize(T + alpha*ca.sumsqr(U))

            opti.solver('ipopt', {'expand': True},
                        {'tol': 10**-8, 'print_level': print_level})
            sol = opti.solve()

            pos_x_opt = sol.value(pos_x)
            pos_y_opt = sol.value(pos_y)
            ang_th_opt = sol.value(ang_th)
            u_opt = sol.value(U)
            T_opt = sol.value(T)
            mprim.append({'x': pos_x_opt, 'y': pos_y_opt, 'th': ang_th_opt,
                          'u': u_opt, 'T': T_opt, 'ds': T_opt*v})
        return np.array(mprim)
コード例 #27
0
    def construct_problem(self):
        """Formulate optimal control problem"""

        dt = self.sample_rate

        # Create an casadi.Opti instance.
        opti = casadi.Opti('conic')

        d0 = opti.parameter()
        th0 = opti.parameter()
        v = opti.parameter()
        curvature = opti.parameter(self.N)

        X = opti.variable(2, self.N + 1)
        proj_error = X[0, :]
        head_error = X[1, :]

        # Control variable (steering angle)
        Delta = opti.variable(self.N)

        # Goal function we wish to minimize
        ### YOUR CODE HERE ###
        # Use the casadi.sumsqr function to compute sum-of-squares
        d_sum = casadi.sumsqr(proj_error)
        th_sum = casadi.sumsqr(head_error)
        u_sum = casadi.sumsqr(Delta)

        J = self.gamma_d * d_sum + self.gamma_theta * th_sum + self.gamma_u * u_sum

        opti.minimize(J)

        # Simulate the system forwards using RK4 and the implemented error model.
        for k in range(self.N):
            k1 = self.error_model(X[:, k], v, Delta[k], curvature[k])
            k2 = self.error_model(X[:, k] + dt / 2 * k1, v, Delta[k],
                                  curvature[k])
            k3 = self.error_model(X[:, k] + dt / 2 * k2, v, Delta[k],
                                  curvature[k])
            k4 = self.error_model(X[:, k] + dt * k3, v, Delta[k], curvature[k])
            x_next = X[:, k] + dt / 6 * (k1 + 2 * k2 + 2 * k3 + k4)
            opti.subject_to(X[:, k + 1] == x_next)

        # Problem constraints.
        opti.subject_to(proj_error[0] == d0)
        opti.subject_to(head_error[0] == th0)
        opti.subject_to(
            opti.bounded(-self.steer_limit, Delta, self.steer_limit))

        # The cost function is quadratic and the problem is linear by design,
        # this is utilized when choosing solver.
        # Other possible solvers provided by CasADi include: 'qpoases' and 'ipopt'...

        opts_dict = {
            "print_iter": False,
            "print_time": 0,
            #"max_iter": 100
        }

        opti.solver('qrqp', opts_dict)

        return opti.to_function('f', [d0, th0, v, curvature], [Delta])
コード例 #28
0
    def solve_opti(self, abs_t, x_0, x_ss, N, last_u, x_guess=None, u_guess=None, expl_constraints=None, verbose=False):
        opti = ca.Opti()

        x = opti.variable(self.n_x, N+1)
        u = opti.variable(self.n_u, N)
        slack = opti.variable(self.n_x)

        if x_guess is not None:
            opti.set_initial(x, x_guess)
        if u_guess is not None:
            opti.set_initial(u, u_guess)
        opti.set_initial(slack, np.zeros(self.n_x))

        da_lim = self.agent.da_lim
        ddf_lim = self.agent.ddf_lim

        opti.subject_to(x[:,0] == np.squeeze(x_0))
        opti.subject_to(opti.bounded(ddf_lim[0]*self.dt, u[0,0]-last_u[0], ddf_lim[1]*self.dt))
        opti.subject_to(opti.bounded(da_lim[0]*self.dt, u[1,0]-last_u[1], da_lim[1]*self.dt))

        stage_cost = 0
        for i in range(N):
            stage_cost = stage_cost + 1

            beta = ca.atan2(self.l_r*ca.tan(u[0,i]), self.l_f+self.l_r)
            opti.subject_to(x[0,i+1] == x[0,i] + self.dt*x[3,i]*ca.cos(x[2,i] + beta))
            opti.subject_to(x[1,i+1] == x[1,i] + self.dt*x[3,i]*ca.sin(x[2,i] + beta))
            opti.subject_to(x[2,i+1] == x[2,i] + self.dt*x[3,i]*ca.sin(beta))
            opti.subject_to(x[3,i+1] == x[3,i] + self.dt*u[1,i])

            if self.F is not None:
                opti.subject_to(ca.mtimes(self.F, x[:,i]) <= self.b)
            if self.H is not None:
                opti.subject_to(ca.mtimes(self.H, u[:,i]) <= self.g)

            if expl_constraints is not None:
                V = expl_constraints[i][0]
                w = expl_constraints[i][1]
                opti.subject_to(ca.mtimes(V, x[:2,i]) + w <= np.zeros(len(w)))

            if i < N-1:
                opti.subject_to(opti.bounded(ddf_lim[0]*self.dt, u[0,i+1]-u[0,i], ddf_lim[1]*self.dt))
                opti.subject_to(opti.bounded(da_lim[0]*self.dt, u[1,i+1]-u[1,i], da_lim[1]*self.dt))

        opti.subject_to(x[:,N] - x_ss == slack)

        slack_cost = 1e6*ca.sumsqr(slack)
        total_cost = stage_cost + slack_cost
        opti.minimize(total_cost)

        solver_opts = {
            "mu_strategy" : "adaptive",
            "mu_init" : 1e-5,
            "mu_min" : 1e-15,
            "barrier_tol_factor" : 1,
            "print_level" : 0,
            "linear_solver" : "ma27"
            }
        # solver_opts = {}
        plugin_opts = {"verbose" : False, "print_time" : False, "print_out" : False}
        opti.solver('ipopt', plugin_opts, solver_opts)

        sol = opti.solve()

        slack_val = sol.value(slack)
        if la.norm(slack_val) > 1e-6:
            print('Warning! Solved, but with slack norm of %g is greater than 1e-8!' % la.norm(slack_val))

        if sol.stats()['success'] and la.norm(slack_val) <= 1e-6:
            print('Solve success, with slack norm of %g!' % la.norm(slack_val))
            feasible = True

            x_pred = sol.value(x)
            u_pred = sol.value(u)
            sol_cost = sol.value(stage_cost)
        else:
            # print(sol.stats()['return_status'])
            # print(opti.debug.show_infeasibilities())
            feasible = False
            x_pred = None
            u_pred = None
            sol_cost = None

            # pdb.set_trace()

        return x_pred, u_pred, sol_cost
コード例 #29
0
    def solve(self, abs_t, x_0, x_ss, N, last_u, x_guess=None, u_guess=None, expl_constraints=None, verbose=False):
        x = ca.SX.sym('x', self.n_x*(N+1))
        u = ca.SX.sym('y', self.n_u*N)
        slack = ca.SX.sym('slack', self.n_x)

        z = ca.vertcat(x, u, slack)

        # Flatten candidate solutions into 1-D array (- x_1 -, - x_2 -, ..., - x_N+1 -)
        if x_guess is not None:
            x_guess_flat = x_guess.flatten(order='F')
        else:
            x_guess_flat = np.zeros(self.n_x*(N+1))

        if u_guess is not None:
            u_guess_flat = u_guess.flatten(order='F')
        else:
            u_guess_flat = np.zeros(self.n_u*N)
        slack_guess = np.zeros(self.n_x)

        z_guess = np.concatenate((x_guess_flat, u_guess_flat, slack_guess))

        lb_slack = [-1000.0]*self.n_x
        ub_slack = [1000.0]*self.n_x

        # Box constraints on decision variables
        lb_x = x_0.tolist() + self.state_lb*(N) + self.input_lb*(N) + lb_slack
        ub_x = x_0.tolist() + self.state_ub*(N) + self.input_ub*(N) + ub_slack

        # Constraints on functions of decision variables
        lb_g = []
        ub_g = []

        stage_cost = 0
        constraint = []
        for i in range(N):
            stage_cost += 1

            # Formulate dynamics equality constraints as inequalities
            beta = ca.atan2(self.l_r*ca.tan(u[self.n_u*i+0]), self.l_f+self.l_r)
            constraint = ca.vertcat(constraint, x[self.n_x*(i+1)+0] - (x[self.n_x*i+0] + self.dt*x[self.n_x*i+3]*ca.cos(x[self.n_x*i+2] + beta)))
            constraint = ca.vertcat(constraint, x[self.n_x*(i+1)+1] - (x[self.n_x*i+1] + self.dt*x[self.n_x*i+3]*ca.sin(x[self.n_x*i+2] + beta)))
            constraint = ca.vertcat(constraint, x[self.n_x*(i+1)+2] - (x[self.n_x*i+2] + self.dt*x[self.n_x*i+3]*ca.sin(beta)))
            constraint = ca.vertcat(constraint, x[self.n_x*(i+1)+3] - (x[self.n_x*i+3] + self.dt*u[self.n_u*i+1]))

            lb_g += [0.0]*self.n_x
            ub_g += [0.0]*self.n_x

            # Steering rate constraints
            if self.ddf_lim is not None:
                if i == 0:
                    # Constrain steering rate with respect to previously applied input
                    constraint = ca.vertcat(constraint, u[self.n_u*(i)+0]-last_u[0])
                if i < N-1:
                    # Constrain steering rate along horizon
                    constraint = ca.vertcat(constraint, u[self.n_u*(i+1)+0]-u[self.n_u*(i)+0])
                lb_g += [self.ddf_lim[0]*self.dt]
                ub_g += [self.ddf_lim[1]*self.dt]

            # Throttle rate constraints
            if self.da_lim is not None:
                if i == 0:
                    # Constrain throttle rate
                    constraint = ca.vertcat(constraint, u[self.n_u*i+1]-last_u[1])
                if i < N-1:
                    constraint = ca.vertcat(constraint, u[self.n_u*(i+1)+1]-u[self.n_u*(i)+1])
                lb_g += [self.da_lim[0]*self.dt]
                ub_g += [self.da_lim[1]*self.dt]

            # Exploration constraints on predicted positions of agent
            if expl_constraints is not None:
                V = expl_constraints[i][0]
                w = expl_constraints[i][1]

                for j in range(V.shape[0]):
                    constraint = ca.vertcat(constraint, V[j,0]*x[self.n_x*i+0] + V[j,1]*x[self.n_x*i+1] + w[j])
                    lb_g += [-1e9]
                    ub_g += [0]

        # Formulate terminal soft equality constraint as inequalities
        constraint = ca.vertcat(constraint, x[self.n_x*N:] - x_ss + slack)
        lb_g += [0.0]*self.n_x
        ub_g += [0.0]*self.n_x

        slack_cost = 1e6*ca.sumsqr(slack)
        total_cost = stage_cost + slack_cost

        opts = {'verbose' : False, 'print_time' : 0,
            'ipopt.print_level' : 5,
            'ipopt.mu_strategy' : 'adaptive',
            'ipopt.mu_init' : 1e-5,
            'ipopt.mu_min' : 1e-15,
            'ipopt.barrier_tol_factor' : 1,
            'ipopt.linear_solver': 'ma27'}
        nlp = {'x' : z, 'f' : total_cost, 'g' : constraint}
        solver = ca.nlpsol('solver', 'ipopt', nlp, opts)

        sol = solver(lbx=lb_x, ubx=ub_x, lbg=lb_g, ubg=ub_g, x0=z_guess)

        pdb.set_trace()

        if solver.stats()['success']:
            if la.norm(slack_val) <= 1e-8:
                print('Solve success for safe set point', x_ss)
                feasible = True

                z_sol = np.array(sol['x'])
                x_pred = z_sol[:self.n_x*(N+1)].reshape((N+1,self.n_x)).T
                u_pred = z_sol[self.n_x*(N+1):self.n_x*(N+1)+self.n_u*N].reshape((N,self.n_u)).T
                slack_val = z_sol[self.n_x*(N+1)+self.n_u*N:]
                cost_val = sol['f']
            else:
                print('Warning! Solved, but with slack norm of %g is greater than 1e-8!' % la.norm(slack_val))
                feasible = False
                x_pred = None
                u_pred = None
                cost_val = None
        else:
            print(solver.stats()['return_status'])
            feasible = False
            x_pred = None
            u_pred = None
            cost_val = None

            # pdb.set_trace()

        # pdb.set_trace()

        return x_pred, u_pred, cost_val
コード例 #30
0
ファイル: falcon9_rtls.py プロジェクト: janismac/yamtof
def create_rocket_stage_phase(mocp, phase_name, p, **kwargs):
    is_powered = kwargs['is_powered']
    has_heating = kwargs['has_heating']
    drag_area = kwargs['drag_area']
    maxQ = kwargs['maxQ']
    init_mass = kwargs['init_mass']
    init_position = kwargs['init_position']
    init_velocity = kwargs['init_velocity']

    if is_powered:
        init_steering = kwargs['init_steering']

        engine_type = kwargs['engine_type']
        engine_count = kwargs['engine_count']
        mdot_max = kwargs['mdot_max']
        mdot_min = kwargs['mdot_min']

        if engine_type == 'vacuum':
            effective_exhaust_velocity = kwargs['effective_exhaust_velocity']
        else:
            exhaust_velocity = kwargs['exhaust_velocity']
            exit_area = kwargs['exit_area']
            exit_pressure = kwargs['exit_pressure']

    duration = mocp.create_phase(phase_name, init=0.001, n_intervals=2)

    # Total vessel mass
    mass = mocp.add_trajectory(phase_name, 'mass', init=init_mass)

    # ECEF position vector
    rx = mocp.add_trajectory(phase_name, 'rx', init=init_position[0])
    ry = mocp.add_trajectory(phase_name, 'ry', init=init_position[1])
    rz = mocp.add_trajectory(phase_name, 'rz', init=init_position[2])
    r_vec = casadi.vertcat(rx, ry, rz)

    # ECEF velocity vector
    vx = mocp.add_trajectory(phase_name, 'vx', init=init_velocity[0])
    vy = mocp.add_trajectory(phase_name, 'vy', init=init_velocity[1])
    vz = mocp.add_trajectory(phase_name, 'vz', init=init_velocity[2])
    v_vec = casadi.vertcat(vx, vy, vz)

    if is_powered:

        # ECEF steering unit vector
        ux = mocp.add_trajectory(phase_name, 'ux', init=init_steering[0])
        uy = mocp.add_trajectory(phase_name, 'uy', init=init_steering[1])
        uz = mocp.add_trajectory(phase_name, 'uz', init=init_steering[2])

        u_vec = casadi.vertcat(ux, uy, uz)

        # Engine mass flow for a SINGLE engine
        mdot = mocp.add_trajectory(phase_name, 'mdot', init=mdot_max)
        mdot_rate = mocp.add_trajectory(phase_name, 'mdot_rate', init=0.0)

    if has_heating:
        heat = mocp.add_trajectory(phase_name, 'heat', init=0.0)

    ## Dynamics

    r_squared = rx**2 + ry**2 + rz**2 + 1e-8
    v_squared = vx**2 + vy**2 + vz**2 + 1e-8
    airspeed = v_squared**0.5
    r = r_squared**0.5
    altitude = r - 1
    r3_inv = r_squared**(-1.5)

    mocp.add_path_output(
        'downrange_distance',
        casadi.asin(
            casadi.norm_2(
                casadi.cross(
                    r_vec / r,
                    casadi.SX(
                        [p.launch_site_x, p.launch_site_y,
                         p.launch_site_z])))))
    mocp.add_path_output('altitude', altitude)
    mocp.add_path_output('airspeed', airspeed)
    mocp.add_path_output('vertical_speed', (r_vec.T / r) @ v_vec)
    mocp.add_path_output(
        'horizontal_speed_ECEF',
        casadi.norm_2(v_vec - ((r_vec.T / r) @ v_vec) * (r_vec / r)))

    # Gravitational acceleration (mu=1 is omitted)
    gx = -r3_inv * rx
    gy = -r3_inv * ry
    gz = -r3_inv * rz

    # Atmosphere
    atmosphere_fraction = casadi.exp(-altitude / p.scale_height)
    air_pressure = p.air_pressure_MSL * atmosphere_fraction
    air_density = p.air_density_MSL * atmosphere_fraction
    dynamic_pressure = 0.5 * air_density * v_squared
    mocp.add_path_output('dynamic_pressure', dynamic_pressure)

    if is_powered:
        lateral_airspeed = casadi.sqrt(
            casadi.sumsqr(v_vec - ((v_vec.T @ u_vec) * u_vec)) + 1e-8)
        lateral_dynamic_pressure = 0.5 * air_density * lateral_airspeed * airspeed  # Same as Q * sin(alpha), but without trigonometry
        mocp.add_path_output('lateral_dynamic_pressure',
                             lateral_dynamic_pressure)

        mocp.add_path_output(
            'alpha',
            casadi.acos((v_vec.T @ u_vec) / airspeed) * 180.0 / pi)
        mocp.add_path_output(
            'pitch', 90.0 - casadi.acos((r_vec.T / r) @ u_vec) * 180.0 / pi)

    # Thrust force
    if is_powered:
        if engine_type == 'vacuum':
            engine_thrust = effective_exhaust_velocity * mdot
        else:
            engine_thrust = exhaust_velocity * mdot + exit_area * (
                exit_pressure - air_pressure)

        total_thrust = engine_thrust * engine_count
        mocp.add_path_output('total_thrust', total_thrust)

        Tx = total_thrust * ux
        Ty = total_thrust * uy
        Tz = total_thrust * uz
    else:
        Tx = 0.0
        Ty = 0.0
        Tz = 0.0

    # Drag force
    drag_factor = (-0.5 * drag_area) * air_density * airspeed
    Dx = drag_factor * vx
    Dy = drag_factor * vy
    Dz = drag_factor * vz

    # Coriolis Acceleration
    ax_Coriolis = 2 * vy * p.body_rotation_speed
    ay_Coriolis = -2 * vx * p.body_rotation_speed
    az_Coriolis = 0.0

    # Centrifugal Acceleration
    ax_Centrifugal = rx * p.body_rotation_speed**2
    ay_Centrifugal = ry * p.body_rotation_speed**2
    az_Centrifugal = 0.0

    # Acceleration
    ax = gx + ax_Coriolis + ax_Centrifugal + (Dx + Tx) / mass
    ay = gy + ay_Coriolis + ay_Centrifugal + (Dy + Ty) / mass
    az = gz + az_Coriolis + az_Centrifugal + (Dz + Tz) / mass

    if is_powered:
        mocp.set_derivative(mass, -engine_count * mdot)
        mocp.set_derivative(mdot, mdot_rate)
    else:
        mocp.set_derivative(mass, 0.0)

    mocp.set_derivative(rx, vx)
    mocp.set_derivative(ry, vy)
    mocp.set_derivative(rz, vz)
    mocp.set_derivative(vx, ax)
    mocp.set_derivative(vy, ay)
    mocp.set_derivative(vz, az)

    # Heat flow
    heat_flow = 1e-4 * (air_density**0.5) * (airspeed**3.15)
    mocp.add_path_output('heat_flow', heat_flow)
    if has_heating:
        mocp.set_derivative(heat, heat_flow)

    ## Constraints

    # Duration is positive and bounded
    mocp.add_constraint(duration > 0.006)
    mocp.add_constraint(duration < 10.0)

    # Mass is positive
    mocp.add_path_constraint(mass > 1e-6)

    # maxQ soft constraint
    weight_dynamic_pressure_penalty = mocp.get_parameter(
        'weight_dynamic_pressure_penalty')
    slack_maxQ = mocp.add_trajectory(phase_name, 'slack_maxQ', init=10.0)
    mocp.add_path_constraint(slack_maxQ > 0.0)
    mocp.add_mean_objective(weight_dynamic_pressure_penalty * slack_maxQ)
    mocp.add_path_constraint(dynamic_pressure / maxQ < 1.0 + slack_maxQ)

    if is_powered:
        # u is a unit vector
        mocp.add_path_constraint(ux**2 + uy**2 + uz**2 == 1.0)

        # Do not point down
        mocp.add_path_constraint((r_vec / r).T @ u_vec > -0.2)

        # Engine flow limits
        mocp.add_path_constraint(mdot_min < mdot)
        mocp.add_path_constraint(mdot < mdot_max)

        # Throttle rate reduction
        mocp.add_mean_objective(1e-6 * mdot_rate**2)

        # Lateral dynamic pressure penalty
        weight_lateral_dynamic_pressure_penalty = mocp.get_parameter(
            'weight_lateral_dynamic_pressure_penalty')
        mocp.add_mean_objective(
            weight_lateral_dynamic_pressure_penalty *
            (lateral_dynamic_pressure / p.maxQ_sin_alpha)**8)

    variables = locals().copy()
    RocketStagePhase = collections.namedtuple('RocketStagePhase',
                                              sorted(list(variables.keys())))
    return RocketStagePhase(**variables)