mpc.cpp

Go to the documentation of this file.

#include <mpc/mpc.h>

#include <Eigen/Core>
#include <Eigen/Eigenvalues>
#include <algorithm>
#include <cmath>

namespace MPC {

        bool DDP_init(Dynamics dyn, DynamicsD dynd, CostFunction cost, CostFunctionD costd) {
                bool res = true;
                res &= setDynamicFunction(dyn, dynd);
                res &= setCostFunction(cost, costd);
                return res;
        }
        
        int forwardPass(const VectorT &x0, const MatrixT &u, const LocalValueFunction &v, const Trajectory &in, Trajectory &out, float alpha) {
     
                float tolDiv = 1e6;
                int n = x0.rows();
                int m = u.rows();
                int N = u.cols();

                out.x.setZero(n, N+1);
                out.x.col(0) = x0;
                out.u.setZero(m, N);
                out.cost.setZero(N+1);

                for (int i = 0; i < N; ++i) {
                        VectorT dx = out.x.col(i) - in.x.col(i);
                        out.u.col(i) = u.col(i) + alpha * v.l.col(i) + v.L[i] * dx;

                        dynamicsAndCost(out.x.col(i), out.u.col(i), out.x.col(i+1), out.cost(i));

                        if (out.x.col(i+1).array().abs().maxCoeff() > tolDiv) {
                                return 1;
                        }
                }

                out.cost(N) = dynamicsAndCost(out.x.col(N), VectorT::Zero(m));
                return 0;
        }

        MatrixT weightedMatrixAvg(const VectorT &a, const std::vector<MatrixT> &B) {
                assert(B.size() > 1);
                assert(a.size() == int(B.size()));
                int n = B[0].rows();
                int m = B[0].cols();
                MatrixT res = MatrixT::Zero(n, m);
                for (size_t i = 0; i < B.size(); ++i) {
                        res += a[i] * B[i];
                }
                return res;
        }


        int backwardPass(const TrajectoryDeriv &td, LocalValueFunction &v, float lam) {
                // get dimensions
                int n = td[0].cx.rows();
                int N = td.size() - 1;
                int m = td[0].cu.rows();
                
                // init the local value function
                v.dV.setZero(2);
                v.l.setZero(m,N);
                v.Vx.setZero(n,N+1);
               
                assert(v.L.size() == N);
                assert(v.Vxx.size() == N+1);
                for (int k = 0; k < N; ++k) {
                        v.L[k].setZero(m,n);
                        v.Vxx[k].setZero(n,n);
                        //v.Vxx.push_back(MatrixT::Zero(n,n));
                }

                // copy last entry
                v.Vx.col(N) = td[N].cx; 
                v.Vxx[N] = td[N].cxx;
                
                int diverge=-1;

                // traverse time backwards
                for (int k = N-1; k >= 0; --k) {
                        /* TODO: remove debug
                           LOG(td[k].cu.rows() << "x" << td[k].cu.cols());
                           LOG(v.Vx.col(k+1).rows() << "x" << v.Vx.col(k+1).cols());
                           LOG(td[k].fu.rows() << "x" << td[k].fu.cols());
                           VectorT t = (v.Vx.col(k+1).transpose() *  td[k].fu);
                           LOG("t " << t.rows() << "x" << t.cols());
                           t += td[k].cu;
                           LOG("t " << t.rows() << "x" << t.cols());
                        */
                        VectorT Qu = (v.Vx.col(k+1).transpose() * td[k].fu);
                        Qu += td[k].cu;
                        VectorT Qx = (v.Vx.col(k+1).transpose() * td[k].fx);
                        Qx += td[k].cx;
                        MatrixT Qxu = td[k].cxu + td[k].fx.transpose() * (v.Vxx[k+1].transpose() * td[k].fu) + weightedMatrixAvg(v.Vx.col(k+1), td[k].fxu);
                        
                        MatrixT Quu = td[k].cuu + td[k].fu.transpose() * (v.Vxx[k+1].transpose() * td[k].fu) + weightedMatrixAvg(v.Vx.col(k+1), td[k].fuu);
                        MatrixT Qxx = td[k].cxx + td[k].fx.transpose() * (v.Vxx[k+1].transpose() * td[k].fx) + weightedMatrixAvg(v.Vx.col(k+1), td[k].fxx);
                        MatrixT VxxF = v.Vxx[k+1] + lam * MatrixT::Identity(n,n);
                        MatrixT QxuF = td[k].cxu + td[k].fx.transpose() * (VxxF * td[k].fu) + weightedMatrixAvg(v.Vx.col(k+1), td[k].fxu);
                        MatrixT QuuF = td[k].cuu + td[k].fu.transpose() * (VxxF * td[k].fu) + weightedMatrixAvg(v.Vx.col(k+1), td[k].fuu);


                        Eigen::EigenSolver<MatrixT> eigen_solver(QuuF);
                        Eigen::VectorXcf eigs = eigen_solver.eigenvalues();
                                
                        if (eigs.imag().maxCoeff() != 0.) {
                                diverge = k;
                                break;
                        }
                                        
                        MatrixT Quu_inv = QuuF.inverse();
                        v.l.col(k) = -1. * Quu_inv * Qu;
                                
                        v.L[k] = -1. * (Quu_inv * QxuF.transpose());

                        // TODO: check whether dot prod or not
                        v.dV(0) = v.dV(0) + v.l.col(k).dot(Qu);

                        v.dV(1) = v.dV(1) + 0.5 * (v.l.col(k).dot(Quu * v.l.col(k)));

                        v.Vx.col(k) = Qx + v.L[k].transpose() * Qu + Qxu * v.l.col(k) + v.L[k].transpose() * (Quu * v.l.col(k));

                        v.Vxx[k] = Qxx + v.L[k].transpose() * Qxu.transpose() + Qxu * v.L[k] + v.L[k].transpose()* (Quu * v.L[k]);
                 
                        v.Vxx[k] = 0.5 * (v.Vxx[k] + v.Vxx[k].transpose());                                                   
                }  

                return diverge;
        }


        bool update_lam_dlam(float &lam, float &dlam, const DDPParams &params) {
                dlam = std::max(dlam * params.lambdaFactor, params.lambdaFactor);
                lam = std::max(lam * dlam, params.lambdaMin);
                if (lam > params.lambdaMax) {
                        LOG("reached lambdaMax curr: " << lam << " max " << params.lambdaMax);
                        return true;
                }
                return false;
        }

        void DDP(const VectorT &x0, const MatrixT &u0, const DDPParams &params, const int maxIter, Trajectory &curr_traj) {
                bool DEBUG = false;
                float lam = params.lamInit;
                float dlam = params.dlamInit;
                MatrixT grad_norm_tmp;
                
                LocalValueFunction v;
                Trajectory next_traj;
                TrajectoryDeriv deriv;

                int n = x0.size();
                int m = u0.rows();
                int N = u0.cols();

                // setup trajectory and v
                curr_traj.x.setZero(n, N);
                curr_traj.u.setZero(m, N);
                curr_traj.cost.setZero(N);

                next_traj.x.setZero(n, N);
                next_traj.u.setZero(m, N);
                next_traj.cost.setZero(N);

                v.Vx.setZero(n, N);
                v.dV.setZero(2);
                v.l.setZero(m, N);
                v.Vxx.resize(N+1);
                v.L.resize(N);

                grad_norm_tmp.setZero(m, N);

                for (int k = 0; k < N; ++k) {
                        v.L[k].setZero(m,n);
                        v.Vxx[k].setZero(n,n);
                        //v.Vxx.push_back(MatrixT::Zero(n,n));
                }
                v.Vxx[N].setZero(n,n);

                // first do a forward simulation of the initial trajectory
                int diverge = forwardPass(x0, u0, v, curr_traj, curr_traj, 1.);

                for (int i = 0; i < maxIter; ++i) {
                        // first step: differentiate dynamics and cost along the current trajectory
                        deriv.resize(curr_traj.x.cols());
                        dynamicsAndCostDerivative(curr_traj, deriv);

                        // second step: backward pass computing optimal control law and cost
                        while (true) {
                                diverge = backwardPass(deriv, v, lam);
                                if (diverge > -1) {
                                        LOG("increasing lambda diverging in step: " << diverge);
                                        if (update_lam_dlam(lam, dlam, params))
                                                break;
                                } else {
                                        // DONE yehaw :)
                                        break;
                                }
                        }

           
                        grad_norm_tmp.array() = v.l.array().abs() / (curr_traj.u.array().abs() + 1.);
                        float g_norm = grad_norm_tmp.colwise().maxCoeff().mean();
                        if (g_norm < params.tolGrad) {
                                LOG("Sucess gradient norm converged!");
                                break;
                        }
                        // third step: forward pass again to find new policy and cost
                        float alpha = 1.;
                        while (true) {
                                diverge = forwardPass(x0, curr_traj.u, v, curr_traj, next_traj, alpha);
                                /* TODO remove debug only
                                LOG("u0 " << u0);
                                LOG("u " << curr_traj.u.transpose());
                                LOG("ccost " << curr_traj.cost.transpose());
                                LOG("ccost " << curr_traj.cost.sum());
                                LOG("ncost " << next_traj.cost.sum());
                                */
                                float dcost = curr_traj.cost.sum() - next_traj.cost.sum() - 1e10 * diverge;
                                float expected = -alpha * (v.dV(0) + alpha * v.dV(1));
                                if (std::fabs(expected) < 1e-8) {
                                        float sign = (expected > 0.) ? 1. : -1.; 
                                        expected =  sign * 1e-8;
                                }

                                float z = 1.;
                                if (std::fabs(dcost - expected) > 1e-4)
                                        z = dcost / expected;
                                LOG("iteration " << i << " cost expected ratio: dcost = " << dcost << " expected = " << expected << " z = " << z);
                                if (params.zMin > z) {
                                        alpha = alpha / 2.;
                                        if (alpha < 1e-4) {
                                                update_lam_dlam(lam, dlam, params);
                                                break;
                                        } else {
                                                continue;
                                        }
                                } else {
                                        curr_traj = next_traj;
                                        dlam = std::min(dlam / params.lambdaFactor, 1.f / params.lambdaFactor);
                                        lam  = lam * dlam * (lam > params.lambdaMin);
                                        break;
                                }
                        }
                        if (lam > params.lambdaMax) {
                                LOG("reached lambdaMax!");
                                break;
                        }
                        if (DEBUG) { 
                                // TODO: possibly add some debug output here
                        }
                }
                
        }
}