y2014/control_loops/python/dt_mpc.py

#!/usr/bin/python3

import numpy
import scipy.optimize
from matplotlib import pylab

dt = 0.05

# This module computes a NMPC which solves the non-holonomic point
# stabilization problem for a mobile robot.  The algorithm is currently
# too slow to run on a robot in real-time, but is fast enough to use offline
# and to use to compare different parameters for NMPCs.
#
# Inital algorithm from http://www.ece.ufrgs.br/~fetter/icma05_608.pdf


def cartesian_to_polar(X_cartesian):
    """Converts a cartesian coordinate to polar.

  Args:
    X_cartesian, numpy.matrix[3, 1] with x, y, theta as rows.

  Returns:
    numpy.matrix[3, 1] with e, phi, alpha as rows.
  """
    phi = numpy.arctan2(X_cartesian[1, 0], X_cartesian[0, 0])
    return numpy.matrix([[numpy.hypot(X_cartesian[0, 0], X_cartesian[1, 0])],
                         [phi], [X_cartesian[2, 0] - phi]])


def polar_to_cartesian(X_polar):
    """Converts a polar coordinate to cartesian.

  Args:
    X_polar, numpy.matrix[3, 1] with e, phi, alpha as rows.

  Returns:
    numpy.matrix[3, 1] with x, y, theta as rows.
  """
    return numpy.matrix([[X_polar[0, 0] * numpy.cos(X_polar[1, 0])],
                         [X_polar[0, 0] * numpy.sin(X_polar[1, 0])],
                         [X_polar[1, 0] + X_polar[2, 0]]])


def simulate_dynamics(X_cartesian, U):
    """Calculates the robot location after 1 timestep.

  Args:
    X_cartesian, numpy.matrix[3, 1] with the starting location in
        cartesian coordinates.
    U, numpy.matrix[2, 1] with velocity, angular_velocity as rows.

  Returns:
    numpy.matrix[3, 1] with the cartesian coordinate.
  """
    X_cartesian += numpy.matrix([[U[0, 0] * numpy.cos(X_cartesian[2, 0]) * dt],
                                 [U[0, 0] * numpy.sin(X_cartesian[2, 0]) * dt],
                                 [U[1, 0] * dt]])

    return X_cartesian


def U_from_array(U_array):
    """Converts the U array from the optimizer to a bunch of column vectors.

  Args:
    U_array, numpy.array[N] The U coordinates in v, av, v, av, ...

  Returns:
    numpy.matrix[2, N/2] with [[v, v, v, ...], [av, av, av, ...]]
  """
    return numpy.matrix(U_array).reshape((2, -1), order='F')


def U_to_array(U_matrix):
    """Converts the U matrix to the U array for the optimizer.

  Args:
    U_matrix, numpy.matrix[2, N/2] with [[v, v, v, ...], [av, av, av, ...]]

  Returns:
    numpy.array[N] The U coordinates in v, av, v, av, ...
  """
    return numpy.array(U_matrix.reshape((1, -1), order='F'))


def cost(U_array, X_cartesian):
    """Computes the cost given the inital position and the U array.

  Args:
    U_array: numpy.array[N] The U coordinates.
    X_cartesian: numpy.matrix[3, 1] The cartesian coordinates of the starting
        location.

  Returns:
    double, The quadratic cost of evaluating U.
  """
    X_cartesian_mod = X_cartesian.copy()
    U_matrix = U_from_array(U_array)
    my_cost = 0
    Q = numpy.matrix([[0.01, 0, 0], [0, 0.01, 0], [0, 0, 0.01]]) / dt / dt
    R = numpy.matrix([[0.001, 0], [0, 0.001]]) / dt / dt
    for U in U_matrix.T:
        # TODO(austin): Let it go to the point from either side.
        U = U.T
        X_cartesian_mod = simulate_dynamics(X_cartesian_mod, U)
        X_current_polar = cartesian_to_polar(X_cartesian_mod)
        my_cost += U.T * R * U + X_current_polar.T * Q * X_current_polar

    return my_cost


if __name__ == '__main__':
    X_cartesian = numpy.matrix([[-1.0], [-1.0], [0.0]])
    x_array = []
    y_array = []
    theta_array = []

    e_array = []
    phi_array = []
    alpha_array = []

    cost_array = []

    time_array = []
    u0_array = []
    u1_array = []

    num_steps = 20

    U_array = numpy.zeros((num_steps * 2))
    for i in range(400):
        print('Iteration', i)
        # Solve the NMPC
        U_array, fx, _, _, _ = scipy.optimize.fmin_slsqp(cost,
                                                         U_array.copy(),
                                                         bounds=[(-1, 1),
                                                                 (-7, 7)] *
                                                         num_steps,
                                                         args=(X_cartesian, ),
                                                         iter=500,
                                                         full_output=True)
        U_matrix = U_from_array(U_array)

        # Simulate the dynamics
        X_cartesian = simulate_dynamics(X_cartesian, U_matrix[:, 0])

        # Save variables for plotting.
        cost_array.append(fx[0, 0])
        u0_array.append(U_matrix[0, 0])
        u1_array.append(U_matrix[1, 0])
        x_array.append(X_cartesian[0, 0])
        y_array.append(X_cartesian[1, 0])
        theta_array.append(X_cartesian[2, 0])
        time_array.append(i * dt)
        X_polar = cartesian_to_polar(X_cartesian)
        e_array.append(X_polar[0, 0])
        phi_array.append(X_polar[1, 0])
        alpha_array.append(X_polar[2, 0])

        U_array = U_to_array(
            numpy.hstack((U_matrix[:, 1:], numpy.matrix([[0], [0]]))))

        if fx < 1e-05:
            print('Cost is', fx, 'after', i, 'cycles, finishing early')
            break

    # Plot
    pylab.figure('xy')
    pylab.plot(x_array, y_array, label='trajectory')

    pylab.figure('time')
    pylab.plot(time_array, x_array, label='x')
    pylab.plot(time_array, y_array, label='y')
    pylab.plot(time_array, theta_array, label='theta')
    pylab.plot(time_array, e_array, label='e')
    pylab.plot(time_array, phi_array, label='phi')
    pylab.plot(time_array, alpha_array, label='alpha')
    pylab.plot(time_array, cost_array, label='cost')
    pylab.legend()

    pylab.figure('u')
    pylab.plot(time_array, u0_array, label='u0')
    pylab.plot(time_array, u1_array, label='u1')
    pylab.legend()

    pylab.show()