Пример #1
0
    def __init__(self,
                 rtree,
                 q_nom,
                 control_period=0.001,
                 print_period=1.0,
                 sim=True,
                 cost_pub=False):

        LeafSystem.__init__(self)

        self.rtree = rtree
        self.nq = rtree.get_num_positions()
        self.nv = rtree.get_num_velocities()
        self.nu = rtree.get_num_actuators()
        self.cost_pub = cost_pub

        dim = 3  # 3D
        nd = 4  # for friction cone approx, hard coded for now
        self.nc = 4  # number of contacts; TODO don't hard code

        self.nf = self.nc * nd  # number of contact force variables
        self.ncf = 2  # number of loop constraint forces; TODO don't hard code
        self.neps = self.nc * dim  # number of slack variables for contact

        if sim:
            self.sim = True
            self._DeclareInputPort(PortDataType.kVectorValued,
                                   self.nq + self.nv)
            self._DeclareDiscreteState(self.nu)
            self._DeclareVectorOutputPort(BasicVector(self.nu),
                                          self.CopyStateOutSim)

            self.print_period = print_period
            self.last_print_time = -print_period

            self.last_v = np.zeros(3)
            self.lcmgl = lcmgl("Balancing-plan", true_lcm.LCM())
            # self.lcmgl = None
        else:
            self.sim = False
            self._DeclareInputPort(PortDataType.kVectorValued, 1)
            self._DeclareInputPort(PortDataType.kVectorValued,
                                   kCassiePositions)
            self._DeclareInputPort(PortDataType.kVectorValued,
                                   kCassieVelocities)
            self._DeclareInputPort(PortDataType.kVectorValued, 2)
            self._DeclareDiscreteState(kCassieActuators * 8 + 1)
            self._DeclareVectorOutputPort(
                BasicVector(1),
                lambda c, o: self.CopyStateOut(c, o, 8 * kCassieActuators, 1))
            self._DeclareVectorOutputPort(
                BasicVector(kCassieActuators), lambda c, o: self.CopyStateOut(
                    c, o, 0, kCassieActuators))  #torque
            self._DeclareVectorOutputPort(
                BasicVector(kCassieActuators), lambda c, o: self.CopyStateOut(
                    c, o, kCassieActuators, kCassieActuators))  #motor_pos
            self._DeclareVectorOutputPort(
                BasicVector(kCassieActuators), lambda c, o: self.CopyStateOut(
                    c, o, 2 * kCassieActuators, kCassieActuators))  #motor_vel
            self._DeclareVectorOutputPort(
                BasicVector(kCassieActuators), lambda c, o: self.CopyStateOut(
                    c, o, 3 * kCassieActuators, kCassieActuators))  #kp
            self._DeclareVectorOutputPort(
                BasicVector(kCassieActuators), lambda c, o: self.CopyStateOut(
                    c, o, 4 * kCassieActuators, kCassieActuators))  #kd
            self._DeclareVectorOutputPort(
                BasicVector(kCassieActuators), lambda c, o: self.CopyStateOut(
                    c, o, 5 * kCassieActuators, kCassieActuators))  #ki
            self._DeclareVectorOutputPort(
                BasicVector(kCassieActuators), lambda c, o: self.CopyStateOut(
                    c, o, 6 * kCassieActuators, kCassieActuators))  #leak
            self._DeclareVectorOutputPort(
                BasicVector(kCassieActuators), lambda c, o: self.CopyStateOut(
                    c, o, 7 * kCassieActuators, kCassieActuators))  #clamp
        self._DeclarePeriodicDiscreteUpdate(0.0005)

        self.q_nom = q_nom
        self.com_des = rtree.centerOfMass(self.rtree.doKinematics(q_nom))
        self.u_des = CassieFixedPointTorque()

        self.lfoot = rtree.FindBody("toe_left")
        self.rfoot = rtree.FindBody("toe_right")
        self.springs = 2 + np.array([
            rtree.FindIndexOfChildBodyOfJoint("knee_spring_left_fixed"),
            rtree.FindIndexOfChildBodyOfJoint("knee_spring_right_fixed"),
            rtree.FindIndexOfChildBodyOfJoint("ankle_spring_joint_left"),
            rtree.FindIndexOfChildBodyOfJoint("ankle_spring_joint_right")
        ])

        umin = np.zeros(self.nu)
        umax = np.zeros(self.nu)
        ii = 0
        for motor in rtree.actuators:
            umin[ii] = motor.effort_limit_min
            umax[ii] = motor.effort_limit_max
            ii += 1

        slack_limit = 10.0
        self.initialized = False

        #------------------------------------------------------------
        # Add Decision Variables ------------------------------------
        prog = MathematicalProgram()
        qddot = prog.NewContinuousVariables(self.nq, "joint acceleration")
        u = prog.NewContinuousVariables(self.nu, "input")
        bar = prog.NewContinuousVariables(self.ncf, "four bar forces")
        beta = prog.NewContinuousVariables(self.nf, "friction forces")
        eps = prog.NewContinuousVariables(self.neps, "slack variable")
        nparams = prog.num_vars()

        #------------------------------------------------------------
        # Problem Constraints ---------------------------------------
        self.con_u_lim = prog.AddBoundingBoxConstraint(umin, umax,
                                                       u).evaluator()
        self.con_fric_lim = prog.AddBoundingBoxConstraint(0, 100000,
                                                          beta).evaluator()
        self.con_slack_lim = prog.AddBoundingBoxConstraint(
            -slack_limit, slack_limit, eps).evaluator()

        bar_con = np.zeros((self.ncf, self.nq))
        self.con_4bar = prog.AddLinearEqualityConstraint(
            bar_con, np.zeros(self.ncf), qddot).evaluator()

        if self.nc > 0:
            dyn_con = np.zeros(
                (self.nq, self.nq + self.nu + self.ncf + self.nf))
            dyn_vars = np.concatenate((qddot, u, bar, beta))
            self.con_dyn = prog.AddLinearEqualityConstraint(
                dyn_con, np.zeros(self.nq), dyn_vars).evaluator()

            foot_con = np.zeros((self.neps, self.nq + self.neps))
            foot_vars = np.concatenate((qddot, eps))
            self.con_foot = prog.AddLinearEqualityConstraint(
                foot_con, np.zeros(self.neps), foot_vars).evaluator()

        else:
            dyn_con = np.zeros((self.nq, self.nq + self.nu + self.ncf))
            dyn_vars = np.concatenate((qddot, u, bar))
            self.con_dyn = prog.AddLinearEqualityConstraint(
                dyn_con, np.zeros(self.nq), dyn_vars).evaluator()

        #------------------------------------------------------------
        # Problem Costs ---------------------------------------------
        self.Kp_com = 50
        self.Kd_com = 2.0 * sqrt(self.Kp_com)
        # self.Kp_qdd = 10*np.eye(self.nq)#np.diag(np.diag(H)/np.max(np.diag(H)))
        self.Kp_qdd = np.diag(
            np.concatenate(([10, 10, 10, 5, 5, 5], np.zeros(self.nq - 6))))
        self.Kd_qdd = 1.0 * np.sqrt(self.Kp_qdd)
        self.Kp_qdd[self.springs, self.springs] = 0
        self.Kd_qdd[self.springs, self.springs] = 0

        com_ddot_des = np.zeros(dim)
        qddot_des = np.zeros(self.nq)

        self.w_com = 0
        self.w_qdd = np.zeros(self.nq)
        self.w_qdd[:6] = 10
        self.w_qdd[8:10] = 1
        self.w_qdd[3:6] = 1000
        # self.w_qdd[self.springs] = 0
        self.w_u = 0.001 * np.ones(self.nu)
        self.w_u[2:4] = 0.01
        self.w_slack = 0.1

        self.cost_qdd = prog.AddQuadraticErrorCost(np.diag(self.w_qdd),
                                                   qddot_des,
                                                   qddot).evaluator()
        self.cost_u = prog.AddQuadraticErrorCost(np.diag(self.w_u), self.u_des,
                                                 u).evaluator()
        self.cost_slack = prog.AddQuadraticErrorCost(
            self.w_slack * np.eye(self.neps), np.zeros(self.neps),
            eps).evaluator()

        # self.cost_com = prog.AddQuadraticCost().evaluator()
        # self.cost_qdd = prog.AddQuadraticCost(
        #     2*np.diag(self.w_qdd), -2*np.matmul(qddot_des, np.diag(self.w_qdd)), qddot).evaluator()
        # self.cost_u = prog.AddQuadraticCost(
        #     2*np.diag(self.w_u), -2*np.matmul(self.u_des, np.diag(self.w_u)) , u).evaluator()
        # self.cost_slack = prog.AddQuadraticCost(
        #     2*self.w_slack*np.eye(self.neps), np.zeros(self.neps), eps).evaluator()

        REG = 1e-8
        self.cost_reg = prog.AddQuadraticErrorCost(
            REG * np.eye(nparams), np.zeros(nparams),
            prog.decision_variables()).evaluator()
        # self.cost_reg = prog.AddQuadraticCost(2*REG*np.eye(nparams),
        #     np.zeros(nparams), prog.decision_variables()).evaluator()

        #------------------------------------------------------------
        # Solver settings -------------------------------------------
        self.solver = GurobiSolver()
        # self.solver = OsqpSolver()
        prog.SetSolverOption(SolverType.kGurobi, "Method", 2)

        #------------------------------------------------------------
        # Save MathProg Variables -----------------------------------
        self.qddot = qddot
        self.u = u
        self.bar = bar
        self.beta = beta
        self.eps = eps
        self.prog = prog
Пример #2
0
    def __init__(self, rtree, q_nom, control_period=0.001):

        self.rtree = rtree
        self.nq = rtree.get_num_positions()
        self.nv = rtree.get_num_velocities()
        self.nu = rtree.get_num_actuators()

        dim = 3  # 3D
        nd = 4  # for friction cone approx, hard coded for now
        self.nc = 4  # number of contacts; TODO don't hard code

        self.nf = self.nc * nd  # number of contact force variables
        self.ncf = 2  # number of loop constraint forces; TODO don't hard code
        self.neps = self.nc * dim  # number of slack variables for contact

        self.q_nom = q_nom
        self.com_des = rtree.centerOfMass(self.rtree.doKinematics(q_nom))
        self.u_des = CassieFixedPointTorque()

        self.lfoot = rtree.FindBody("toe_left")
        self.rfoot = rtree.FindBody("toe_right")
        self.springs = 2 + np.array([
            rtree.FindIndexOfChildBodyOfJoint("knee_spring_left_fixed"),
            rtree.FindIndexOfChildBodyOfJoint("knee_spring_right_fixed"),
            rtree.FindIndexOfChildBodyOfJoint("ankle_spring_joint_left"),
            rtree.FindIndexOfChildBodyOfJoint("ankle_spring_joint_right")
        ])

        umin = np.zeros(self.nu)
        umax = np.zeros(self.nu)
        ii = 0
        for motor in rtree.actuators:
            umin[ii] = motor.effort_limit_min
            umax[ii] = motor.effort_limit_max
            ii += 1

        slack_limit = 10.0
        self.initialized = False

        #------------------------------------------------------------
        # Add Decision Variables ------------------------------------
        prog = MathematicalProgram()
        qddot = prog.NewContinuousVariables(self.nq, "joint acceleration")
        u = prog.NewContinuousVariables(self.nu, "input")
        bar = prog.NewContinuousVariables(self.ncf, "four bar forces")
        beta = prog.NewContinuousVariables(self.nf, "friction forces")
        eps = prog.NewContinuousVariables(self.neps, "slack variable")
        nparams = prog.num_vars()

        #------------------------------------------------------------
        # Problem Constraints ---------------------------------------
        self.con_u_lim = prog.AddBoundingBoxConstraint(umin, umax,
                                                       u).evaluator()
        self.con_fric_lim = prog.AddBoundingBoxConstraint(0, 100000,
                                                          beta).evaluator()
        self.con_slack_lim = prog.AddBoundingBoxConstraint(
            -slack_limit, slack_limit, eps).evaluator()

        bar_con = np.zeros((self.ncf, self.nq))
        self.con_4bar = prog.AddLinearEqualityConstraint(
            bar_con, np.zeros(self.ncf), qddot).evaluator()

        if self.nc > 0:
            dyn_con = np.zeros(
                (self.nq, self.nq + self.nu + self.ncf + self.nf))
            dyn_vars = np.concatenate((qddot, u, bar, beta))
            self.con_dyn = prog.AddLinearEqualityConstraint(
                dyn_con, np.zeros(self.nq), dyn_vars).evaluator()

            foot_con = np.zeros((self.neps, self.nq + self.neps))
            foot_vars = np.concatenate((qddot, eps))
            self.con_foot = prog.AddLinearEqualityConstraint(
                foot_con, np.zeros(self.neps), foot_vars).evaluator()

        else:
            dyn_con = np.zeros((self.nq, self.nq + self.nu + self.ncf))
            dyn_vars = np.concatenate((qddot, u, bar))
            self.con_dyn = prog.AddLinearEqualityConstraint(
                dyn_con, np.zeros(self.nq), dyn_vars).evaluator()

        #------------------------------------------------------------
        # Problem Costs ---------------------------------------------
        self.Kp_com = 50
        self.Kd_com = 2.0 * sqrt(self.Kp_com)
        # self.Kp_qdd = 10*np.eye(self.nq)#np.diag(np.diag(H)/np.max(np.diag(H)))
        self.Kp_qdd = np.diag(
            np.concatenate(([10, 10, 10, 5, 5, 5], np.zeros(self.nq - 6))))
        self.Kd_qdd = 1.0 * np.sqrt(self.Kp_qdd)
        self.Kp_qdd[self.springs, self.springs] = 0
        self.Kd_qdd[self.springs, self.springs] = 0

        com_ddot_des = np.zeros(dim)
        qddot_des = np.zeros(self.nq)

        self.w_com = 0
        self.w_qdd = np.zeros(self.nq)
        self.w_qdd[:6] = 10
        self.w_qdd[8:10] = 1
        self.w_qdd[3:6] = 1000
        # self.w_qdd[self.springs] = 0
        self.w_u = 0.001 * np.ones(self.nu)
        self.w_u[2:4] = 0.01
        self.w_slack = 0.1

        self.cost_qdd = prog.AddQuadraticErrorCost(np.diag(self.w_qdd),
                                                   qddot_des,
                                                   qddot).evaluator()
        self.cost_u = prog.AddQuadraticErrorCost(np.diag(self.w_u), self.u_des,
                                                 u).evaluator()
        self.cost_slack = prog.AddQuadraticErrorCost(
            self.w_slack * np.eye(self.neps), np.zeros(self.neps),
            eps).evaluator()

        # self.cost_com = prog.AddQuadraticCost().evaluator()
        # self.cost_qdd = prog.AddQuadraticCost(
        #     2*np.diag(self.w_qdd), -2*np.matmul(qddot_des, np.diag(self.w_qdd)), qddot).evaluator()
        # self.cost_u = prog.AddQuadraticCost(
        #     2*np.diag(self.w_u), -2*np.matmul(self.u_des, np.diag(self.w_u)) , u).evaluator()
        # self.cost_slack = prog.AddQuadraticCost(
        #     2*self.w_slack*np.eye(self.neps), np.zeros(self.neps), eps).evaluator()

        REG = 1e-8
        self.cost_reg = prog.AddQuadraticErrorCost(
            REG * np.eye(nparams), np.zeros(nparams),
            prog.decision_variables()).evaluator()
        # self.cost_reg = prog.AddQuadraticCost(2*REG*np.eye(nparams),
        #     np.zeros(nparams), prog.decision_variables()).evaluator()

        #------------------------------------------------------------
        # Solver settings -------------------------------------------
        self.solver = GurobiSolver()
        # self.solver = OsqpSolver()
        prog.SetSolverOption(SolverType.kGurobi, "Method", 2)

        #------------------------------------------------------------
        # Save MathProg Variables -----------------------------------
        self.qddot = qddot
        self.u = u
        self.bar = bar
        self.beta = beta
        self.eps = eps
        self.prog = prog
Пример #3
0
def solveOptimization(state_init,
                      t_impact,
                      impact_combination,
                      T,
                      u_guess=None,
                      x_guess=None,
                      h_guess=None):
    prog = MathematicalProgram()
    h = prog.NewContinuousVariables(T, name='h')
    u = prog.NewContinuousVariables(rows=T + 1,
                                    cols=2 * n_quadrotors,
                                    name='u')
    x = prog.NewContinuousVariables(rows=T + 1,
                                    cols=6 * n_quadrotors + 4 * n_balls,
                                    name='x')
    dv = prog.decision_variables()

    prog.AddBoundingBoxConstraint([h_min] * T, [h_max] * T, h)

    for i in range(n_quadrotors):
        sys = Quadrotor2D()
        context = sys.CreateDefaultContext()
        dir_coll_constr = DirectCollocationConstraint(sys, context)
        ind_x = 6 * i
        ind_u = 2 * i

        for t in range(T):
            impact_indices = impact_combination[np.argmax(
                np.abs(t - t_impact) <= 1)]
            quad_ind, ball_ind = impact_indices[0], impact_indices[1]

            if quad_ind == i and np.any(
                    t == t_impact
            ):  # Don't add Direct collocation constraint at impact
                continue
            elif quad_ind == i and (np.any(t == t_impact - 1)
                                    or np.any(t == t_impact + 1)):
                prog.AddConstraint(
                    eq(
                        x[t + 1,
                          ind_x:ind_x + 3], x[t, ind_x:ind_x + 3] + h[t] *
                        x[t + 1, ind_x + 3:ind_x + 6]))  # Backward euler
                prog.AddConstraint(
                    eq(x[t + 1, ind_x + 3:ind_x + 6], x[t,
                                                        ind_x + 3:ind_x + 6])
                )  # Zero-acceleration assumption during this time step. Should maybe replace with something less naive
            else:
                AddDirectCollocationConstraint(
                    dir_coll_constr, np.array([[h[t]]]),
                    x[t, ind_x:ind_x + 6].reshape(-1, 1),
                    x[t + 1, ind_x:ind_x + 6].reshape(-1, 1),
                    u[t, ind_u:ind_u + 2].reshape(-1, 1),
                    u[t + 1, ind_u:ind_u + 2].reshape(-1, 1), prog)

    for i in range(n_balls):
        sys = Ball2D()
        context = sys.CreateDefaultContext()
        dir_coll_constr = DirectCollocationConstraint(sys, context)
        ind_x = 6 * n_quadrotors + 4 * i

        for t in range(T):
            impact_indices = impact_combination[np.argmax(
                np.abs(t - t_impact) <= 1)]
            quad_ind, ball_ind = impact_indices[0], impact_indices[1]

            if ball_ind == i and np.any(
                    t == t_impact
            ):  # Don't add Direct collocation constraint at impact
                continue
            elif ball_ind == i and (np.any(t == t_impact - 1)
                                    or np.any(t == t_impact + 1)):
                prog.AddConstraint(
                    eq(
                        x[t + 1,
                          ind_x:ind_x + 2], x[t, ind_x:ind_x + 2] + h[t] *
                        x[t + 1, ind_x + 2:ind_x + 4]))  # Backward euler
                prog.AddConstraint(
                    eq(x[t + 1,
                         ind_x + 2:ind_x + 4], x[t, ind_x + 2:ind_x + 4] +
                       h[t] * np.array([0, -9.81])))
            else:
                AddDirectCollocationConstraint(
                    dir_coll_constr, np.array([[h[t]]]),
                    x[t, ind_x:ind_x + 4].reshape(-1, 1),
                    x[t + 1, ind_x:ind_x + 4].reshape(-1, 1),
                    u[t, 0:0].reshape(-1, 1), u[t + 1, 0:0].reshape(-1,
                                                                    1), prog)

    # Initial conditions
    prog.AddLinearConstraint(eq(x[0, :], state_init))

    # Final conditions
    prog.AddLinearConstraint(eq(x[T, 0:14], state_final[0:14]))
    # Quadrotor final conditions (full state)
    for i in range(n_quadrotors):
        ind = 6 * i
        prog.AddLinearConstraint(
            eq(x[T, ind:ind + 6], state_final[ind:ind + 6]))

    # Ball final conditions (position only)
    for i in range(n_balls):
        ind = 6 * n_quadrotors + 4 * i
        prog.AddLinearConstraint(
            eq(x[T, ind:ind + 2], state_final[ind:ind + 2]))

    # Input constraints
    for i in range(n_quadrotors):
        prog.AddLinearConstraint(ge(u[:, 2 * i], -20.0))
        prog.AddLinearConstraint(le(u[:, 2 * i], 20.0))
        prog.AddLinearConstraint(ge(u[:, 2 * i + 1], -20.0))
        prog.AddLinearConstraint(le(u[:, 2 * i + 1], 20.0))

    # Don't allow quadrotor to pitch more than 60 degrees
    for i in range(n_quadrotors):
        prog.AddLinearConstraint(ge(x[:, 6 * i + 2], -np.pi / 3))
        prog.AddLinearConstraint(le(x[:, 6 * i + 2], np.pi / 3))

    # Ball position constraints
    # for i in range(n_balls):
    #     ind_i = 6*n_quadrotors + 4*i
    #     prog.AddLinearConstraint(ge(x[:,ind_i],-2.0))
    #     prog.AddLinearConstraint(le(x[:,ind_i], 2.0))
    #     prog.AddLinearConstraint(ge(x[:,ind_i+1],-3.0))
    #     prog.AddLinearConstraint(le(x[:,ind_i+1], 3.0))

    # Impact constraint
    quad_temp = Quadrotor2D()

    for i in range(n_quadrotors):
        for j in range(n_balls):
            ind_q = 6 * i
            ind_b = 6 * n_quadrotors + 4 * j
            for t in range(T):
                if np.any(
                        t == t_impact
                ):  # If quad i and ball j impact at time t, add impact constraint
                    impact_indices = impact_combination[np.argmax(
                        t == t_impact)]
                    quad_ind, ball_ind = impact_indices[0], impact_indices[1]
                    if quad_ind == i and ball_ind == j:
                        # At impact, witness function == 0
                        prog.AddConstraint(lambda a: np.array([
                            CalcClosestDistanceQuadBall(a[0:3], a[3:5])
                        ]).reshape(1, 1),
                                           lb=np.zeros((1, 1)),
                                           ub=np.zeros((1, 1)),
                                           vars=np.concatenate(
                                               (x[t, ind_q:ind_q + 3],
                                                x[t,
                                                  ind_b:ind_b + 2])).reshape(
                                                      -1, 1))
                        # At impact, enforce discrete collision update for both ball and quadrotor
                        prog.AddConstraint(
                            CalcPostCollisionStateQuadBallResidual,
                            lb=np.zeros((6, 1)),
                            ub=np.zeros((6, 1)),
                            vars=np.concatenate(
                                (x[t, ind_q:ind_q + 6], x[t, ind_b:ind_b + 4],
                                 x[t + 1, ind_q:ind_q + 6])).reshape(-1, 1))
                        prog.AddConstraint(
                            CalcPostCollisionStateBallQuadResidual,
                            lb=np.zeros((4, 1)),
                            ub=np.zeros((4, 1)),
                            vars=np.concatenate(
                                (x[t, ind_q:ind_q + 6], x[t, ind_b:ind_b + 4],
                                 x[t + 1, ind_b:ind_b + 4])).reshape(-1, 1))

                        # rough constraints to enforce hitting center-ish of paddle
                        prog.AddLinearConstraint(
                            x[t, ind_q] - x[t, ind_b] >= -0.01)
                        prog.AddLinearConstraint(
                            x[t, ind_q] - x[t, ind_b] <= 0.01)
                        continue
                # Everywhere else, witness function must be > 0
                prog.AddConstraint(lambda a: np.array([
                    CalcClosestDistanceQuadBall(a[ind_q:ind_q + 3], a[
                        ind_b:ind_b + 2])
                ]).reshape(1, 1),
                                   lb=np.zeros((1, 1)),
                                   ub=np.inf * np.ones((1, 1)),
                                   vars=x[t, :].reshape(-1, 1))

    # Don't allow quadrotor collisions
    # for t in range(T):
    #     for i in range(n_quadrotors):
    #         for j in range(i+1, n_quadrotors):
    #             prog.AddConstraint((x[t,6*i]-x[t,6*j])**2 + (x[t,6*i+1]-x[t,6*j+1])**2 >= 0.65**2)

    # Quadrotors stay on their own side
    # prog.AddLinearConstraint(ge(x[:, 0], 0.3))
    # prog.AddLinearConstraint(le(x[:, 6], -0.3))

    ###############################################################################
    # Set up initial guesses
    initial_guess = np.empty(prog.num_vars())

    # # initial guess for the time step
    prog.SetDecisionVariableValueInVector(h, h_guess, initial_guess)

    x_init[0, :] = state_init
    prog.SetDecisionVariableValueInVector(x, x_guess, initial_guess)

    prog.SetDecisionVariableValueInVector(u, u_guess, initial_guess)

    solver = SnoptSolver()
    print("Solving...")
    result = solver.Solve(prog, initial_guess)

    # print(GetInfeasibleConstraints(prog,result))
    # be sure that the solution is optimal
    assert result.is_success()

    print(f'Solution found? {result.is_success()}.')

    #################################################################################
    # Extract results
    # get optimal solution
    h_opt = result.GetSolution(h)
    x_opt = result.GetSolution(x)
    u_opt = result.GetSolution(u)
    time_breaks_opt = np.array([sum(h_opt[:t]) for t in range(T + 1)])
    u_opt_poly = PiecewisePolynomial.ZeroOrderHold(time_breaks_opt, u_opt.T)
    # x_opt_poly = PiecewisePolynomial.Cubic(time_breaks_opt, x_opt.T, False)
    x_opt_poly = PiecewisePolynomial.FirstOrderHold(
        time_breaks_opt, x_opt.T
    )  # Switch to first order hold instead of cubic because cubic was taking too long to create
    #################################################################################
    # Create list of K matrices for time varying LQR
    context = quad_plant.CreateDefaultContext()
    breaks = copy.copy(
        time_breaks_opt)  #np.linspace(0, x_opt_poly.end_time(), 100)

    K_samples = np.zeros((breaks.size, 12 * n_quadrotors))

    for i in range(n_quadrotors):
        K = None
        for j in range(breaks.size):
            context.SetContinuousState(
                x_opt_poly.value(breaks[j])[6 * i:6 * (i + 1)])
            context.FixInputPort(
                0,
                u_opt_poly.value(breaks[j])[2 * i:2 * (i + 1)])
            linear_system = FirstOrderTaylorApproximation(quad_plant, context)
            A = linear_system.A()
            B = linear_system.B()
            try:
                K, _, _ = control.lqr(A, B, Q, R)
            except:
                assert K is not None, "Failed to calculate initial K for quadrotor " + str(
                    i)
                print("Warning: Failed to calculate K at timestep", j,
                      "for quadrotor", i, ". Using K from previous timestep")

            K_samples[j, 12 * i:12 * (i + 1)] = K.reshape(-1)

    K_poly = PiecewisePolynomial.ZeroOrderHold(breaks, K_samples.T)

    return u_opt_poly, x_opt_poly, K_poly, h_opt