#!/usr/bin/env python
# coding: utf-8

# # Optimal Control and Parameter Identification of Dynamcal Systems with Direct Collocation and SymPy
# 
# Jason K. Moore   
# July 8, 2015  
# SciPy 2015, Austin, Texas, USA
# 
# <table style="border: 0;">
# <tr style="border: 0;">
# <td style="border: 0;">
# <img width=300px src="NewCSU-stacked.svg"/>
# </td>
# <td style="border: 0;">
# <img width=100px src="sympy-logo.svg" />
# </td>
# <td style="border: 0;">
# <img width=100px src="pydy-logo.svg" />
# </td>
# </tr>
# </table>

# # Motivation: Identification of Gait Control
# 
# <img src="mocap-to-control.jpg" />

# # Previous Use Cases: Optimal Gait on Earth, Moon, Mars
# 
# > Ackermann, M. and van den Bogert,  A. J. Predictive simulation of gait at low gravity reveals skipping as the preferred
# locomotion strategy. Journal of Biomechanics, 45, 2012
# 
# <table>
#   <tr>
#     <td>
#      <video width="320" height="240" controls autoplay loop>
#        <source src="walking-videos/earth.mp4" type="video/mp4">
#        <source src="walking-videos/earth.ogg" type="video/ogg">
#      </video>
#      <p>Earth: g=9.8 m/s<sup>-2</sup></p>
#     </td>
#         <td>
#      <video width="320" height="240" controls autoplay loop>
#        <source src="walking-videos/moon.mp4" type="video/mp4">
#        <source src="walking-videos/moon.ogg" type="video/ogg">
#      </video>
#      <p>Moon: g=1.6 m/s<sup>-2</sup></p>
#     </td>
#   </tr>
# </table>

# # Trajectory Optimization Example
# 
# <img src="pendulum.svg" style="float: right;" width=200px />
# 
# ## 1-DoF, 1 parameter pendulum equations of motion
# 
# $\dot{\mathbf{x}} = \begin{bmatrix} \dot{\theta}(t) \\ \dot{\omega}(t) \end{bmatrix} = \begin{bmatrix} \omega(t) \\ \frac{g}{l} \sin{\theta}(t) + \mathbf{T} \end{bmatrix}$
# 
# ## Objective: Minimize "effort" and bound the states at
# 
# $J(\mathbf{T}) = \min\limits_{\mathbf{T}} \int_{t_0}^{t_f} \mathbf{T}^2 dt$
# 
# ## Boundary conditions
# 
# - $\mathbf{\theta}(t_0) = 0$
# - $\mathbf{\theta}(t_f) = \pi$
# 
# > Pendulum figure: https://commons.wikimedia.org/wiki/File:Pendulum_gravity.svg, Creative Commons Attribution-Share Alike 3.0 Unported

# # Parameter Identification
# 
# Find the parameters $\mathbf{p}$ such that the difference between the model simulation, $\mathbf{y}$, and measurements, $\mathbf{y}_m$ is minimized.
# 
# ### Dynamic system
# 
# - Equations of Motion: $\dot{\mathbf{x}} = \mathbf{f}(\mathbf{x}, \mathbf{p})$
# - Measurement variables: $\mathbf{y} = \mathbf{g}(\mathbf{x}, \mathbf{p})$
# 
# ### Objective
# 
# $$\min\limits_\mathbf{p} J(\mathbf{p})$$
# where
# $$J(\mathbf{p}) = \int [\mathbf{y}_m - \mathbf{y}(\mathbf{p})]^2 dt$$

# # Optimal Control By Shooting
# 
# - Repeated simulations are computationally costly (i.e. stiff systems are especially slow)
# - Maybe limited to parameterized functions (e.g. polynomials, piecewise constants)
# - Too many free variables
# - Systems may be unstable and thus have an ill-defined objective
# - Local minima are inevitable
#   - May requires a superb guess
#   - May need time intensive global optimization methods

# # Local Minima Example: Simple pendulum
# 
# > Vyasarayani, Chandrika P., Thomas Uchida, Ashwin Carvalho, and John McPhee.
# "Parameter Identification in Dynamic Systems Using the Homotopy Optimization
# Approach". Multibody System Dynamics 26, no. 4 (2011): 411-24.
# 
# <img style="float:right;" src="pendulum-objective.png" width=400px>
# 
# ## 1-DoF, 1 parameter pendulum equations of motion
# 
# $\dot{\mathbf{x}} = \begin{bmatrix} \dot{\theta}(t) \\ \dot{\omega}(t) \end{bmatrix} = \begin{bmatrix} \omega(t) \\ -p \sin{\theta}(t) \end{bmatrix}$
# 
# ## Objective: Minimize least squares
# 
# $J(p) = \min\limits_{p} \int_{t_0}^{t_f} [\theta_m(t) - \theta(\mathbf{x}, p, t)]^2 dt$

# # Direct Collocation
# 
# 
# ## Benefits
# 
# <ul style="display: inline;">
#   <img src="betts-book.jpg" style="float:right;" width=200px />
#   <li>Fast computation times</li>
#   <li>Handles unstable systems with ease</li>
#   <li>Less susceptible to local minima</li>
# </ul>
# 
# ## Disadvantages
# 
# <ul style="display: inline;">
#   <li>Accurate solution requires large number of nodes</li>
#   <li>Memory management for large sparse matrices and operations</li>
#   <li>Tedious and error prone to form gradients, Jacobians, and Hessians</li>
# </ul>

# # Our Direct Collocation Implementation
# 
# ### Implicit Continous Equations of Motion
# 
# No need to solve for $\dot{\mathbf{x}}$.
# 
# $$\mathbf{f}(\dot{\mathbf{x}}, \mathbf{x}, \mathbf{r}, \mathbf{p}, t) = 0$$
# 
# - $\mathbf{x}, \dot{\mathbf{x}}$: states and their derivatives
# - $\mathbf{r}$: exogenous inputs
# - $\mathbf{p}$: constant parameters

# # Discretization
# 
# ### First order discrete integration options:
# 
# Backward Euler:
# 
# $\mathbf{f}(\frac{\mathbf{x}_i - \mathbf{x}_{i-1}}{h} , \mathbf{x}_i, \mathbf{r}_i, \mathbf{p}, t_i) = 0$
# 
# Midpoint Rule:
# 
# $\mathbf{f}(\frac{\mathbf{x}_{i + 1} - \mathbf{x}_i}{h}, \frac{\mathbf{x}_i + \mathbf{x}_{i + 1}}{2}, \frac{\mathbf{r}_i + \mathbf{r}_{i + 1}}{2}, \mathbf{p}, t_i) = 0$
# 
# ### Nonlinear Programming Formulation
#  
# $$\min\limits_{\theta} J(\theta)$$
# 
# $$\textrm{where } \mathbf{\theta} = [\mathbf{x}_i, \dots, \mathbf{x}_N, \mathbf{r}_i, \ldots, \mathbf{r}_N,  \mathbf{p}]$$
# 
# $$\textrm{subject to } \mathbf{f}(\mathbf{x}_i, \mathbf{x}_{i+1}, \mathbf{r}_i, \mathbf{r}_{i+1}, \mathbf{p}, t_i) = 0 \textrm{ and } \theta_L \leq \theta \leq \theta_U$$

# # Software Tool: opty
# 
# ## http://github.com/csu-hmc/opty
# 
# - User specifies continous symbolic:
#   - objective
#   - equations of motion (explicit or implicit)
#   - additional constraints
#   - bounds on free variables: $\mathbf{x}, \mathbf{r}, \mathbf{p}$
# - EoMs can be generated with PyDy (http://pydy.org)
# - Effficient just-in-time compiled C code is generated for functions that evaluate:
#   - objective and its gradient
#   - constraints and its Jacobian
# - NLP problem automatically formed for IPOPT
# - Open source: BSD license
# - Written in Python

# # Example Code: Specify Equations of Motion
# 
# <br />
# 
# $$\dot{\mathbf{x}} = \begin{bmatrix} \dot{\theta}(t) \\ \dot{\omega}(t) \end{bmatrix} = \begin{bmatrix} \omega(t) \\ -p \sin{\theta}(t) + T(t) \end{bmatrix}$$
# 
# ### Symbolic EoM
# 
# ```python
# # Specify symbols for the parameters
# p, t = symbols('p, t')
# 
# # Specify the functions of time
# theta, omega, theta_m, T = symbols('theta, omega, theta_m, T', cls=Function)
# 
# # Specify the symbolic equations of motion
# eom = Matrix([theta(t), omega(t)]).diff(t) - Matrix([omega(t), -p * sin(theta(t) + T(t))])
# ```
# 
# ### Discretization
# 
# ```python
# # Choose discretization values
# num_nodes = 1000
# interval = 0.01  # seconds
# ```

# # Example Code: Trajectory Optimization
# 
# ### Objective
# 
# ```python
# obj = Integral(T(t)**2, t)
# ```
# 
# ### Boundary Contraints
# ```python
# boundary_constraints = (theta(0.0), theta(duration) - pi / 2, omega(0.0), omega(duration))
# ```
# 
# ### Define Problem
# 
# ```python
# prob = Problem(obj,  # symbolic objective
#                eom,  # symbolic equations of motion
#                (theta(t), omega(t)),  # system symbolic states
#                num_nodes,
#                interval,
#                known_parameter_map={p: 9.81},
#                instance_constraints=boundary_constraints)
# ```

# # Example Code: Single Parameter Identification
# 
# ### Objective
# 
# ```python
# obj = Integral((theta_m(t) - theta(t))**2, t)
# ```
# 
# ### Form Problem
# 
# ```python
# prob = Problem(obj,
#                eom,
#                (theta(t), omega(t)),
#                num_nodes,
#                interval,
#                known_trajectory_map={T(t): 0.0, y1_m(t): measured_data},
#                integration_method='midpoint')
# ```

# # Example Code: Solve
# ```python
# # Set an initial guess
# initial_guess = random(prob.num_free)
# 
# # Solve the system
# solution, info = prob.solve(initial_guess)
# ```

# # The Problem Class: Objective
# 
# 1. Converts objective integral to discrete form
# 2. Computes analytic gradient of objective wrt to discrete variables
# 3. Generates efficient NumPy functions for objective and gradient
# 
# *Still in development

# # The Problem Class: Constraints
# 
# 1. Converts continous EoM to discrete versions
# 2. Computes the analytic Jacobian wrt to discrete variables
# 3. Generates efficient C code to "ufuncify" discrete EoM and Jacobian evaluations
# 4. Generates Cython wrapper for C code
# 5. Compiles Cython wrapper
# 6. NumPy functions are generated for instance constraints

# ### The Problem Class: Solve
# 
# 1. IPOPT data structure is formed with cyipopt
# 2. Bounds on variables are set
# 3. IPOPT settings are set
# 4. Solve is called!

# # Computational Speed
# 
# ### Example Larger System
# <img style="float:right;margin=5px;" src="n-pendulum-with-cart.svg" width=400px />
# <ul style="display: inline;">
#   <li>10 link pendulum on sliding cart (not stiff)</li>
#   <li>11 DoF, 22 states, 22 parameters</li>
#   <li>12800 mathematical operations in constraint expressions</li>
#   <li>100 s sampled @ 100 hz</li>
# </ul>

# # Computational Speed
# 
# ### Discretization Variables
# 
# - 10,000 collocation nodes
# - 219,978 constraints
# - 14,518,548 nonzero entries in the Jacobian
# - 220,022 free variables
# 
# ### Timings
# 
# - Integrating with ODEPACK lsoda: **5.6 s**
# - Constraint evaluation: **33 ms (0.033 s)**
# - Jacobian evaluation: **128 ms (0.128 s)**

# # Case Study: Minamal Effort Double Pendulum Swing Up
# 
# 
# <img src="two-link-pendulum-with-cart.svg"style="float:right;margin=5px;" width=400px/>
# 
# <ul style="display: inline;">
#   <li>Two-link pendulum with lateral force applied to base</li> 
#   <li>States: $\mathbf{x}=[q_0 \quad q_1 \quad q_2 \quad u_0 \quad u_1 \quad u_2]^T$</li>
#   <li>Unknown exogoneous input: $\mathbf{F}$</li>
#   <li>Known parameters: $[m_0, m_1, m_2, l_0, l_1]$</li>
#   <li>Boundary conditions:
#     <ul>
#       <li>$q_0(t_0)=0$</li>
#       <li>$q_1(t_0),q_1(t_0)=-\frac{\pi}{2}$</li>
#       <li>$q_1(t_f),q_1(t_f)=\frac{\pi}{2}$</li>
#       <li>$u_0(t_0), u_2(t_0), u_2(t_0), u_0(t_f), u_1(t_f), u_2(t_f)=0$</li>
#     </ul>
#   </li>
# </ul>

# # Trajectory Opimization Problem Specification
# 
# Given the boundary conditions can we find the optimal input trajectory such that the "effort" is minimized?
# 
# $$\min\limits_\theta J(\mathbf{\theta}), \quad
# J(\mathbf{\theta})= \sum_{i=1}^N h [\mathbf{T}_i]^2$$
# 
# where
# 
# $$ \mathbf{\theta} = [\mathbf{x}_1, \ldots, \mathbf{x}_N, \mathbf{T}_i, \ldots \mathbf{T}_N] $$
# 
# Subject to the constraints:
# 
# $$\mathbf{f}(\mathbf{x}_i, \mathbf{T}_i)=0, \quad i=1 \ldots N$$
# 
# And the initial guess:
# 
# $$\mathbf{\theta}_0 = [\mathbf{0}]$$
# 
# For, $N$ = 500:
# 
# - 24008 free variables
# - Jacobian matrix with 38933 non-zero entries

# # Demo
# 
# <img src="two-link-pendulum-with-cart.svg" width=600px />

# # Case Study: Human Control Parameter Identification
# 
# <div>
#   <h2>Plant</h2>
#   <img style="float:right;margin=5px;" src="free-body-diagram.svg" width=200px>
#   <ul style="display: inline;">
#     <li>Torque driven two-link inverted pendulum with an accelerating base.</li> 
#     <li>States: $\mathbf{x}=[\theta_a \quad \theta_h \quad \omega_a \quad \omega_h]^T$</li>
#     <li>Exogoneous inputs:
#       <ul>
#         <li>Controlled: $\mathbf{r}_c = [T_a \quad T_h]^T$</li>
#         <li>Specified: $\mathbf{r}_k = [a]$</li>
#       </ul>
#     </li>
#     <li>Known parameters: $\mathbf{p}_k$</li>
#   </ul>
#   <h3>Open Loop Equations of Motion</h3>
#   <p>$$\dot{\mathbf{x}} = \mathbf{f}_o(\mathbf{x}, \mathbf{r}_c, \mathbf{r}_k, \mathbf{p}_k, t)$$</p>
# </div>

# # Lumped Passive+Active Controller
# 
# - True human controller is practically impossible to isolate and identify
# - Identify a controller for a similar system that causes the same behavior as the real system
# 
# ## Simple State Feedback
# 
# $$\mathbf{r}_c(t) = -\mathbf{K}\mathbf{x}(t)$$
# 
# ## Unknown Parameters
# 
# $$\mathbf{p}_u = \mathrm{vec}(\mathbf{K})$$
# 
# ## Closed Loop Equations of Motion
#   
# $$\dot{\mathbf{x}} = \mathbf{f}_c(\mathbf{x}, \mathbf{r}_k, \mathbf{p}_k, \mathbf{p}_u, t)$$

# # Generate Data
# 
# ## Specify the psuedo-random platform acceleration
#  
# $$a(t)=\sum_{i=1}^{12} A_i\sin(\omega_i t) \quad \mathrm{where,} \quad 0.15 \mathrm{rad/s} < \omega_i < 15.0 \mathrm{rad/s}$$
# 
# ### Choose a stable controller
# 
# $$
# \mathbf{K} =
# \begin{bmatrix}
#   950 & 175 & 185 & 50 \\
#   45 &  290 & 60 & 26
# \end{bmatrix}
# $$

# # Generate Data
# 
# ### Simulate closed loop system under the influence of perturbations for 60 seconds, sampled at 100 Hz
#  
# $$ \dot{\mathbf{x}} = \mathbf{f}_c(\mathbf{x}, \mathbf{r}_k, \mathbf{p}_k, \mathbf{p}_u, t) $$
# 
# ### Add Gaussian measurement noise
# 
# $$\mathbf{x}_m(t) = \mathbf{x}(t) + \mathbf{v}_x(t) \\ a_m(t) = a(t) + v_a(t)$$
# 
# - $\sigma_\theta$ = 0.3 deg
# - $\sigma_\omega$ = 4 deg/s
# - $\sigma_a$ = 0.42 ms<sup>-2</sup>

# # Parameter Identification Problem Specification
# 
# Given noisy measurements of the states, $\mathbf{x}_m$, and the platform acceleration, $a_m$, can we identify the controller parameters $\mathbf{K}$?
# 
# $$\min\limits_\theta J(\mathbf{\theta}), \quad
# J(\mathbf{\theta})= \sum_{i=1}^N h [\mathbf{x}_{mi} - \mathbf{x}_i]^2$$
# 
# where
# 
# $$ \mathbf{\theta} = [\mathbf{x}_1, \ldots, \mathbf{x}_N, \mathbf{p}_u] $$
# 
# Subject to the constraints:
# 
# $$\mathbf{f}_{ci}(\mathbf{x}_i, a_{mi}, \mathbf{p}_u)=0, \quad i=1 \ldots N$$
# 
# And the initial guess:
# 
# $$\mathbf{\theta}_0 = [\mathbf{x}_{m1}, \ldots, \mathbf{x}_{mN}, \mathbf{0}]$$
# 
# For, $N$ = 6000:
# 
# - 24008 free variables
# - 23996 x 24008 Jacobian matrix with 384000 non-zero entries

# # Demo
# 
# <img src="free-body-diagram.svg" width=200px>

# # Results
# 
# Converges in **11 iterations** in **2.8 seconds** of computation time.
# 
# <table border="1" class="dataframe">
#   <thead>
#     <tr style="text-align: right;">
#       <th></th>
#       <th>Known</th>
#       <th>Identified</th>
#       <th>Error</th>
#     </tr>
#   </thead>
#   <tbody>
#     <tr>
#       <th>$k_{00}$</th>
#       <td>950</td>
#       <td>946</td>
#       <td>-0.4%</td>
#     </tr>
#     <tr>
#       <th>$k_{01}$</th>
#       <td>175</td>
#       <td>177</td>
#       <td>1.4%</td>
#     </tr>
#     <tr>
#       <th>$k_{02}$</th>
#       <td>185</td>
#       <td>185</td>
#       <td>-0.2%</td>
#     </tr>
#     <tr>
#       <th>$k_{03}$</th>
#       <td>50</td>
#       <td>55</td>
#       <td>9.4%</td>
#     </tr>
#     <tr>
#       <th>$k_{10}$</th>
#       <td>45</td>
#       <td>45</td>
#       <td>1.1%</td>
#     </tr>
#     <tr>
#       <th>$k_{11}$</th>
#       <td>290</td>
#       <td>289</td>
#       <td>-0.3%</td>
#     </tr>
#     <tr>
#       <th>$k_{12}$</th>
#       <td>60</td>
#       <td>59</td>
#       <td>-2.1%</td>
#     </tr>
#     <tr>
#       <th>$k_{13}$</th>
#       <td>26</td>
#       <td>27</td>
#       <td>4.2%</td>
#     </tr>
#   </tbody>
# </table>
# 

# # Identified State Trajectories
# 
# ![](trajectory-comparison.png)

# # Conclusion
# 
# - Direct collocation is suitable for biomechanical parameter identification and trajectory optimization
# - Computation speeds are orders of magnitude faster than shooting
# - Parameter identification accuracy improves with # nodes
# - Complex problems can be solved with few lines of code and high level mathematical abstractions

# # Questions?
# 
# - Social media: moorepants
# - Personal website: http://moorepants.info
# - SymPy: http://sympy.org
# - PyDy: http://pydy.org
# - opty: http://github.com/csu-hmc/opty
# - Human Motion and Control Lab: http://hmc.csuohio.edu