Continuous-Time Trajectory Optimization

Contents

5. Continuous-Time Trajectory Optimization#

5.1. State-Space Models#

We extend our focus from the discrete-time setting to trajectory optimization in continuous time. Such models are omnipresent in various branches of science and engineering, where the dynamics of physical, biological, or economic systems are often described in terms of continuous-time differential equations. Here, we consider models given by ordinary differential equations (ODEs). However, continuous-time optimal control methods also exist beyond ODEs; for example, using stochastic differential equations (SDEs) or partial differential equations (PDEs).

An example of such state-space representation is the following:

\[\begin{align*} \dot{\mathbf{x}}(t) &= \mathbf{f}(\mathbf{x}(t), \mathbf{u}(t)) \\ \mathbf{y}(t) &= \mathbf{h}(\mathbf{x}(t), \mathbf{u}(t)) \end{align*}\]

The function \(\mathbf{f}\) describes the dynamics of the system, giving the rate of change of the state vector as a function of the current state \(\mathbf{x}(t)\) and control input \(\mathbf{u}(t)\). The function \(\mathbf{h}\) is the output mapping, which determines the measured output \(\mathbf{y}(t)\) based on the current state \(\mathbf{x}(t)\) and control input \(\mathbf{u}(t)\). This state-space representation is reminiscent of recurrent neural networks (RNNs), albeit in discrete time, where we maintain a hidden state and then map it to an observable space.

A state space model is called a linear state space model (or simply a linear system) if the functions \(\mathbf{f}\) and \(\mathbf{h}\) are linear in \(\mathbf{x}(t)\) and \(\mathbf{u}(t)\). In this case, the model can be written as:

\[\begin{align*} \dot{\mathbf{x}}(t) &= \mathbf{A}\mathbf{x}(t) + \mathbf{B}\mathbf{u}(t) \\ \mathbf{y}(t) &= \mathbf{C}\mathbf{x}(t) + \mathbf{D}\mathbf{u}(t) \end{align*}\]

where \(\mathbf{A}\), \(\mathbf{B}\), \(\mathbf{C}\), and \(\mathbf{D}\) are constant matrices. The matrix \(\mathbf{A}\) is called the dynamics matrix, \(\mathbf{B}\) is the control matrix, \(\mathbf{C}\) is the sensor matrix, and \(\mathbf{D}\) is the direct term matrix. If the model does not have a direct term, it means that the control input \(\mathbf{u}(t)\) does not directly influence the output \(\mathbf{y}(t)\).

It is worth noting that linear models like the one presented above are becoming increasingly popular thanks to the development of structured state space models (S4 and such) [Gu et al., 2022]. These models leverage the inherent structure and properties of linear systems to design more efficient and interpretable neural networks for processing sequential data.

5.2. Numerical Methods for Solving ODEs#

An ODE is an implicit representation of a state-space trajectory: it tells us how the state changes in time but not precisely what the state is at any given time. To find out this information, we need to either solve the ODE analytically (for some special structure) or, as we’re going to do, solve them numerically. This numerical procedure is meant to solve what is called an IVP (initial value problem) of the form:

\[ \text{Find } x(t) \text{ given } \dot{x}(t) = f(x(t), t) \text{ and } x(t_0) = x_0 \]

5.2.1. Euler’s Method#

The algorithm to solve this problem is, in its simplest form, a for loop which closely resembles the updates encountered in gradient descent (in fact, gradient descent can be derived from the gradient flow ODE, but that’s another discussion). The so-called explicit Euler’s method can be implemented as follow:

Algorithm 5.1 (Euler’s method)

Input: \(f(x, t)\), \(x_0\), \(t_0\), \(t_{end}\), \(h\)

Output: Approximate solution \(x(t)\) at discrete time points

  1. Initialize \(t = t_0\), \(x = x_0\)

  2. While \(t < t_{end}\):

  3. Compute \(x_{new} = x + h f(x, t)\)

  4. Update \(t = t + h\)

  5. Update \(x = x_{new}\)

  6. Store or output the pair \((t, x)\)

  7. End While

Consider the following simple dynamical system of a ballistic motion model, neglecting air resistance. The state of the system is described by two variables: \(y(t)\): vertical position at time \(t\) and \(v(t)\), the vertical velocity at time \(t\). The corresponding ODE is:

\[\begin{split} \begin{aligned} \frac{dy}{dt} &= v \\ \frac{dv}{dt} &= -g \end{aligned} \end{split}\]

where \(g \approx 9.81 \text{ m/s}^2\) is the acceleration due to gravity. In our code, we use the initial conditions \(y(0) = 0 \text{ m}\) and \(v(0) = v_0 \text{ m/s}\) where \(v_0\) is the initial velocity (in this case, \(v_0 = 20 \text{ m/s}\)). The analytical solution to this system is:

\[\begin{split} \begin{aligned} y(t) &= v_0t - \frac{1}{2}gt^2 \\ v(t) &= v_0 - gt \end{aligned} \end{split}\]

This system models the vertical motion of an object launched upward, reaching a maximum height before falling back down due to gravity.

Euler’s method can be obtained by taking the first-order Taylor expansion of \(x(t)\) at \(t\):

\[ x(t + h) \approx x(t) + h \frac{dx}{dt}(t) = x(t) + h f(x(t), t) \]

Each step of the algorithm therefore involves approximating the function with a linear function of slope \(f\) over the given interval \(h\).

Hide code cell source
import numpy as np
import matplotlib.pyplot as plt

def f(y, t):
    """
    Derivative function for vertical motion under gravity.
    y[0] is position, y[1] is velocity.
    """
    g = 9.81  # acceleration due to gravity (m/s^2)
    return np.array([y[1], -g])

def euler_method(f, y0, t0, t_end, h):
    """
    Implement Euler's method for the entire time range.
    """
    t = np.arange(t0, t_end + h, h)
    y = np.zeros((len(t), 2))
    y[0] = y0
    for i in range(1, len(t)):
        y[i] = y[i-1] + h * f(y[i-1], t[i-1])
    return t, y

def true_solution(t):
    """
    Analytical solution for the ballistic trajectory.
    """
    y0, v0 = 0, 20  # initial height and velocity
    g = 9.81
    return y0 + v0*t - 0.5*g*t**2, v0 - g*t

# Set up the problem
t0, t_end = 0, 4
y0 = np.array([0, 20])  # initial height = 0, initial velocity = 20 m/s

# Different step sizes
step_sizes = [1.0, 0.5, 0.1]
colors = ['r', 'g', 'b']
markers = ['o', 's', '^']

# True solution
t_fine = np.linspace(t0, t_end, 1000)
y_true, v_true = true_solution(t_fine)

# Plotting
plt.figure(figsize=(12, 8))

# Plot Euler approximations
for h, color, marker in zip(step_sizes, colors, markers):
    t, y = euler_method(f, y0, t0, t_end, h)
    plt.plot(t, y[:, 0], color=color, marker=marker, linestyle='--', 
             label=f'Euler h = {h}', markersize=6, markerfacecolor='none')

# Plot true solution last so it's on top
plt.plot(t_fine, y_true, 'k-', label='True trajectory', linewidth=2, zorder=10)

plt.xlabel('Time (s)', fontsize=12)
plt.ylabel('Height (m)', fontsize=12)
plt.title("Euler's Method: Effect of Step Size on Ballistic Trajectory Approximation", fontsize=14)
plt.legend(fontsize=10)
plt.grid(True, linestyle=':', alpha=0.7)

# Add text to explain the effect of step size
plt.text(2.5, 15, "Smaller step sizes\nyield better approximations", 
         bbox=dict(facecolor='white', edgecolor='black', alpha=0.7),
         fontsize=10, ha='center', va='center')

plt.tight_layout()
plt.show()
_images/2b4468a0fd4926f3062b22032921d012df295ab6dbc37f0e70ecb8d59401b7fb.png

Another way to understand Euler’s method is through the fundamental theorem of calculus:

\[ x(t + h) = x(t) + \int_t^{t+h} f(x(\tau), \tau) d\tau \]

We then approximate the integral term with a box of width \(h\) and height \(f\), and therefore of area \(h f\).

Hide code cell source
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.patches import Rectangle

def v(t):
    """
    Velocity function for the ballistic trajectory.
    """
    v0 = 20   # initial velocity (m/s)
    g = 9.81  # acceleration due to gravity (m/s^2)
    return v0 - g * t

def position(t):
    """
    Position function (integral of velocity).
    """
    v0 = 20
    g = 9.81
    return v0*t - 0.5*g*t**2

# Set up the problem
t0, t_end = 0, 2
num_points = 1000
t = np.linspace(t0, t_end, num_points)

# Calculate true velocity and position
v_true = v(t)
x_true = position(t)

# Euler's method with a large step size for visualization
h = 0.5
t_euler = np.arange(t0, t_end + h, h)
x_euler = np.zeros_like(t_euler)

for i in range(1, len(t_euler)):
    x_euler[i] = x_euler[i-1] + h * v(t_euler[i-1])

# Plotting
fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(12, 12), sharex=True)

# Plot velocity function and its approximation
ax1.plot(t, v_true, 'b-', label='True velocity')
ax1.fill_between(t, 0, v_true, alpha=0.3, label='True area (displacement)')

# Add rectangles with hashed pattern, ruler-like annotations, and area values
for i in range(len(t_euler) - 1):
    t_i = t_euler[i]
    v_i = v(t_i)
    rect = Rectangle((t_i, 0), h, v_i, 
                     fill=True, facecolor='red', edgecolor='r', 
                     alpha=0.15, hatch='///')
    ax1.add_patch(rect)
    
    # Add ruler-like annotations
    # Vertical ruler (height)
    ax1.annotate('', xy=(t_i, 0), xytext=(t_i, v_i),
                 arrowprops=dict(arrowstyle='<->', color='red'))
    ax1.text(t_i - 0.05, v_i/2, f'v(t{i}) = {v_i:.2f}', rotation=90, 
             va='center', ha='right', color='red', fontweight='bold')
    
    # Horizontal ruler (width)
    ax1.annotate('', xy=(t_i, -1), xytext=(t_i + h, -1),
                 arrowprops=dict(arrowstyle='<->', color='red'))
    ax1.text(t_i + h/2, -2, f'h = {h}', ha='center', va='top', 
             color='red', fontweight='bold')
    
    # Add area value in the middle of each rectangle
    area = h * v_i
    ax1.text(t_i + h/2, v_i/2, f'Area = {area:.2f}', ha='center', va='center', 
             color='white', fontweight='bold', bbox=dict(facecolor='red', edgecolor='none', alpha=0.7))

# Plot only the points for Euler's method
ax1.plot(t_euler, v(t_euler), 'ro', markersize=6, label="Euler's points")
ax1.set_ylabel('Velocity (m/s)', fontsize=12)
ax1.set_title("Velocity Function and Euler's Approximation", fontsize=14)
ax1.legend(fontsize=10)
ax1.grid(True, linestyle=':', alpha=0.7)
ax1.set_ylim(bottom=-3)  # Extend y-axis to show horizontal rulers

# Plot position function and its approximation
ax2.plot(t, x_true, 'b-', label='True position')
ax2.plot(t_euler, x_euler, 'ro--', label="Euler's approximation", markersize=6, linewidth=2)

# Add vertical arrows and horizontal lines to show displacement and time step
for i in range(1, len(t_euler)):
    t_i = t_euler[i]
    x_prev = x_euler[i-1]
    x_curr = x_euler[i]
    
    # Vertical line for displacement
    ax2.plot([t_i, t_i], [x_prev, x_curr], 'g:', linewidth=2)
    
    # Horizontal line for time step
    ax2.plot([t_i - h, t_i], [x_prev, x_prev], 'g:', linewidth=2)
    
    # Add text to show the displacement value
    displacement = x_curr - x_prev
    ax2.text(t_i + 0.05, (x_prev + x_curr)/2, f'+{displacement:.2f}', 
             color='green', fontweight='bold', va='center')
    
    # Add text to show the time step
    ax2.text(t_i - h/2, x_prev - 0.5, f'h = {h}', 
             color='green', fontweight='bold', ha='center', va='top')

ax2.set_xlabel('Time (s)', fontsize=12)
ax2.set_ylabel('Position (m)', fontsize=12)
ax2.set_title("Position: True vs Euler's Approximation", fontsize=14)
ax2.legend(fontsize=10)
ax2.grid(True, linestyle=':', alpha=0.7)

# Add explanatory text
ax1.text(1.845, 15, "Red hashed areas show\nEuler's approximation\nof the area under the curve", 
         bbox=dict(facecolor='white', edgecolor='black', alpha=0.7),
         fontsize=10, ha='center', va='center')

plt.tight_layout()
plt.show()
_images/178984062060703c45be293899eca6349abad2c019d53fc39a7f4856704b896a.png

5.2.2. Implicit Euler’s Method#

An alternative approach is the Implicit Euler method, also known as the Backward Euler method. Instead of using the derivative at the current point to step forward, it uses the derivative at the end of the interval. This leads to the following update rule:

\[ x_{new} = x + h f(x_{new}, t_{new}) \]

Note that \(x_{new}\) appears on both sides of the equation, making this an implicit method. The algorithm for the Implicit Euler method can be described as follows:

Algorithm 5.2 (Implicit Euler’s Method)

Input: \(f(x, t)\), \(x_0\), \(t_0\), \(t_{end}\), \(h\)

Output: Approximate solution \(x(t)\) at discrete time points

  1. Initialize \(t = t_0\), \(x = x_0\)

  2. While \(t < t_{end}\):

  3. Set \(t_{new} = t + h\)

  4. Solve for \(x_{new}\) in the equation: \(x_{new} = x + h f(x_{new}, t_{new})\)

  5. Update \(t = t_{new}\)

  6. Update \(x = x_{new}\)

  7. Store or output the pair \((t, x)\)

  8. End While

The key difference in the Implicit Euler method is step 4, where we need to solve a (potentially nonlinear) equation to find \(x_{new}\). This is typically done using iterative methods such as fixed-point iteration or Newton’s method.

5.2.2.1. Stiff ODEs#

While the Implicit Euler method requires more computation per step, it often allows for larger step sizes and can provide better stability for certain types of problems, especially stiff ODEs.

Stiff ODEs are differential equations for which certain numerical methods for solving the equation are numerically unstable, unless the step size is taken to be extremely small. These ODEs typically involve multiple processes occurring at widely different rates. In a stiff problem, the fastest-changing component of the solution can make the numerical method unstable unless the step size is extremely small. However, such a small step size may lead to an impractical amount of computation to traverse the entire interval of interest.

For example, consider a chemical reaction where some reactions occur very quickly while others occur much more slowly. The fast reactions quickly approach their equilibrium, but small perturbations in the slower reactions can cause rapid changes in the fast reactions.

A classic example of a stiff ODE is the Van der Pol oscillator with a large parameter. The Van der Pol equation is:

\[ \frac{d^2x}{dt^2} - \mu(1-x^2)\frac{dx}{dt} + x = 0 \]

where \(\mu\) is a scalar parameter. This second-order ODE can be transformed into a system of first-order ODEs by introducing a new variable \(y = \frac{dx}{dt}\):

\[\begin{split} \begin{aligned} \frac{dx}{dt} &= y \\ \frac{dy}{dt} &= \mu(1-x^2)y - x \end{aligned} \end{split}\]

When \(\mu\) is large (e.g., \(\mu = 1000\)), this system becomes stiff. The large \(\mu\) causes rapid changes in \(y\) when \(x\) is near ±1, but slower changes elsewhere. This leads to a solution with sharp transitions followed by periods of gradual change.

5.2.3. Trapezoid Method#

The trapezoid method, also known as the trapezoidal rule, offers improved accuracy and stability compared to the simple Euler method. The name “trapezoid method” comes from the idea of using a trapezoid to approximate the integral term in the fundamental theorem of calculus. This leads to the following update rule:

\[ x_{new} = x + \frac{h}{2}[f(x, t) + f(x_{new}, t_{new})] \]

where \( t_{new} = t + h \). Note that this formula involves \( x_{new} \) on both sides of the equation, making it an implicit method, similar to the implicit Euler method discussed earlier.

Hide code cell source
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.patches import Rectangle, Polygon

def v(t):
    """Velocity function for the ballistic trajectory."""
    v0 = 20   # initial velocity (m/s)
    g = 9.81  # acceleration due to gravity (m/s^2)
    return v0 - g * t

def position(t):
    """Position function (integral of velocity)."""
    v0 = 20
    g = 9.81
    return v0*t - 0.5*g*t**2

# Set up the problem
t0, t_end = 0, 2
num_points = 1000
t = np.linspace(t0, t_end, num_points)

# Calculate true velocity and position
v_true = v(t)
x_true = position(t)

# Euler's method and Trapezoid method with a large step size for visualization
h = 0.5
t_numeric = np.arange(t0, t_end + h, h)
x_euler = np.zeros_like(t_numeric)
x_trapezoid = np.zeros_like(t_numeric)

for i in range(1, len(t_numeric)):
    # Euler's method
    x_euler[i] = x_euler[i-1] + h * v(t_numeric[i-1])
    
    # Trapezoid method (implicit, so we use a simple fixed-point iteration)
    x_trapezoid[i] = x_trapezoid[i-1]
    for _ in range(5):  # 5 iterations should be enough for this simple problem
        x_trapezoid[i] = x_trapezoid[i-1] + h/2 * (v(t_numeric[i-1]) + v(t_numeric[i]))

# Plotting
fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(12, 16), sharex=True)

# Plot velocity function and its approximations
ax1.plot(t, v_true, 'b-', label='True velocity')
ax1.fill_between(t, 0, v_true, alpha=0.3, label='True area (displacement)')

# Add trapezoids and rectangles
for i in range(len(t_numeric) - 1):
    t_i, t_next = t_numeric[i], t_numeric[i+1]
    v_i, v_next = v(t_i), v(t_next)
    
    # Euler's rectangle (hashed pattern)
    rect = Rectangle((t_i, 0), h, v_i, fill=True, facecolor='red', edgecolor='r', alpha=0.15, hatch='///')
    ax1.add_patch(rect)
    
    # Trapezoid (dot pattern)
    trapezoid = Polygon([(t_i, 0), (t_i, v_i), (t_next, v_next), (t_next, 0)], 
                        fill=True, facecolor='green', edgecolor='g', alpha=0.15, hatch='....')
    ax1.add_patch(trapezoid)
    
    # Add area values
    euler_area = h * v_i
    trapezoid_area = h * (v_i + v_next) / 2
    ax1.text(t_i + h/2, v_i/2, f'Euler: {euler_area:.2f}', ha='center', va='bottom', color='red', fontweight='bold')
    ax1.text(t_i + h/2, (v_i + v_next)/4, f'Trapezoid: {trapezoid_area:.2f}', ha='center', va='top', color='green', fontweight='bold')

# Plot points for Euler's and Trapezoid methods
ax1.plot(t_numeric, v(t_numeric), 'ro', markersize=6, label="Euler's points")
ax1.plot(t_numeric, v(t_numeric), 'go', markersize=6, label="Trapezoid points")

ax1.set_ylabel('Velocity (m/s)', fontsize=12)
ax1.set_title("Velocity Function: True vs Numerical Approximations", fontsize=14)
ax1.legend(fontsize=10)
ax1.grid(True, linestyle=':', alpha=0.7)

# Plot position function and its approximations
ax2.plot(t, x_true, 'b-', label='True position')
ax2.plot(t_numeric, x_euler, 'ro--', label="Euler's approximation", markersize=6, linewidth=2)
ax2.plot(t_numeric, x_trapezoid, 'go--', label="Trapezoid approximation", markersize=6, linewidth=2)

ax2.set_xlabel('Time (s)', fontsize=12)
ax2.set_ylabel('Position (m)', fontsize=12)
ax2.set_title("Position: True vs Numerical Approximations", fontsize=14)
ax2.legend(fontsize=10)
ax2.grid(True, linestyle=':', alpha=0.7)

# Add explanatory text
ax1.text(1.76, 17, "Red hashed areas: Euler's approximation\nGreen dotted areas: Trapezoid approximation", 
         bbox=dict(facecolor='white', edgecolor='black', alpha=0.7),
         fontsize=10, ha='center', va='center')

plt.tight_layout()
plt.show()
_images/fb77100ff10ff350a89ed36fd76f9d86205ef8c0db7924c852bf2435b8646013.png

Algorithmically, the trapezoid method can be described as follows:

Algorithm 5.3 (Trapezoid Method)

Input: \( f(x, t) \), \( x_0 \), \( t_0 \), \( t_{end} \), \( h \)

Output: Approximate solution \( x(t) \) at discrete time points

  1. Initialize \( t = t_0 \), \( x = x_0 \)

  2. While \( t < t_{end} \):

  3.   Set \( t_{new} = t + h \)

  4.   Solve for \( x_{new} \)in the equation: \( x_{new} = x + \frac{h}{2}[f(x, t) + f(x_{new}, t_{new})] \)

  5.   Update \( t = t_{new} \)

  6.   Update \( x = x_{new} \)

  7.   Store or output the pair \( (t, x) \)

The trapezoid method can also be derived by averaging the forward Euler and backward Euler methods. Recall that:

  1. Forward Euler method:

    \[ x_{n+1} = x_n + h f(x_n, t_n) \]
  2. Backward Euler method:

    \[ x_{n+1} = x_n + h f(x_{n+1}, t_{n+1}) \]

Taking the average of these two methods yields:

\[\begin{split} \begin{aligned} x_{n+1} &= \frac{1}{2} \left( x_n + h f(x_n, t_n) \right) + \frac{1}{2} \left( x_n + h f(x_{n+1}, t_{n+1}) \right) \\ &= x_n + \frac{h}{2} \left( f(x_n, t_n) + f(x_{n+1}, t_{n+1}) \right) \end{aligned} \end{split}\]

This is precisely the update rule for the trapezoid method. Recall that the forward Euler method approximates the solution by extrapolating linearly using the slope at the beginning of the interval \([t_n, t_{n+1}] \). In contrast, the backward Euler method extrapolates linearly using the slope at the end of the interval. The trapezoid method, on the other hand, averages these two slopes. This averaging provides better approximation properties than either of the methods alone, offering both stability and accuracy. Note finally that unlike the forward or backward Euler methods, the trapezoid method is also symmetric in time. This means that if you were to reverse time and apply the method backward, you would get the same results (up to numerical precision).

5.2.4. Trapezoidal Predictor-Corrector#

The trapezoid method can also be implemented under the so-called predictor-corrector framework. This interpretation reformulates the implicit trapezoid rule into an explicit two-step process:

  1. Predictor Step:
    We make an initial guess for \( x_{n+1} \) using the forward Euler method:

    \[ x_{n+1}^* = x_n + h f(x_n, t_n) \]

    This is our “predictor” step, where \( x_{n+1}^* \) is the predicted value of \( x_{n+1} \).

  2. Corrector Step:
    We then use this predicted value to estimate \( f(x_{n+1}^*, t_{n+1}) \) and apply the trapezoid formula:

    \[ x_{n+1} = x_n + \frac{h}{2} \left[ f(x_n, t_n) + f(x_{n+1}^*, t_{n+1}) \right] \]

    This is our “corrector” step, where the initial guess \( x_{n+1}^* \) is corrected by taking into account the slope at \( (x_{n+1}^*, t_{n+1}) \).

This two-step process is similar to performing one iteration of Newton’s method to solve the implicit trapezoid equation, starting from the Euler prediction. However, to fully solve the implicit equation, multiple iterations would be necessary until convergence is achieved.

Hide code cell source
import numpy as np
import matplotlib.pyplot as plt

def f(y, t):
    """
    Derivative function for vertical motion under gravity.
    y[0] is position, y[1] is velocity.
    """
    g = 9.81  # acceleration due to gravity (m/s^2)
    return np.array([y[1], -g])

def true_solution(t):
    """
    Analytical solution for the ballistic trajectory.
    """
    y0, v0 = 0, 20  # initial height and velocity
    g = 9.81
    return y0 + v0*t - 0.5*g*t**2, v0 - g*t

def trapezoid_method_visual(f, y0, t0, t_end, h):
    """
    Implement the trapezoid method for the entire time range.
    Returns predictor and corrector steps for visualization.
    """
    t = np.arange(t0, t_end + h, h)
    y = np.zeros((len(t), 2))
    y_predictor = np.zeros((len(t), 2))
    y[0] = y_predictor[0] = y0
    for i in range(1, len(t)):
        # Predictor step (Euler forward)
        slope_start = f(y[i-1], t[i-1])
        y_predictor[i] = y[i-1] + h * slope_start
        
        # Corrector step
        slope_end = f(y_predictor[i], t[i])
        y[i] = y[i-1] + h * 0.5 * (slope_start + slope_end)
    
    return t, y, y_predictor

# Set up the problem
t0, t_end = 0, 2
y0 = np.array([0, 20])  # initial height = 0, initial velocity = 20 m/s
h = 0.5  # Step size

# Compute trapezoid method steps
t, y_corrector, y_predictor = trapezoid_method_visual(f, y0, t0, t_end, h)

# Plotting
plt.figure(figsize=(12, 8))

# Plot the true solution for comparison
t_fine = np.linspace(t0, t_end, 1000)
y_true, v_true = true_solution(t_fine)
plt.plot(t_fine, y_true, 'k-', label='True trajectory', linewidth=1.5)

# Plot the predictor and corrector steps
for i in range(len(t)-1):
    # Points for the predictor step
    p0 = [t[i], y_corrector[i, 0]]
    p1_predictor = [t[i+1], y_predictor[i+1, 0]]
    
    # Points for the corrector step
    p1_corrector = [t[i+1], y_corrector[i+1, 0]]
    
    # Plot predictor step
    plt.plot([p0[0], p1_predictor[0]], [p0[1], p1_predictor[1]], 'r--', linewidth=2)
    plt.plot(p1_predictor[0], p1_predictor[1], 'ro', markersize=8)
    
    # Plot corrector step
    plt.plot([p0[0], p1_corrector[0]], [p0[1], p1_corrector[1]], 'g--', linewidth=2)
    plt.plot(p1_corrector[0], p1_corrector[1], 'go', markersize=8)
    
    # Add arrows to show the predictor and corrector adjustments
    plt.arrow(p0[0], p0[1], h, y_predictor[i+1, 0] - p0[1], color='r', width=0.005, 
              head_width=0.02, head_length=0.02, length_includes_head=True, zorder=5)
    plt.arrow(p1_predictor[0], p1_predictor[1], 0, y_corrector[i+1, 0] - y_predictor[i+1, 0], 
              color='g', width=0.005, head_width=0.02, head_length=0.02, length_includes_head=True, zorder=5)

# Add legend entries for predictor and corrector steps
plt.plot([], [], 'r--', label='Predictor step (Forward Euler)')
plt.plot([], [], 'g-', label='Corrector step (Trapezoid)')

# Labels and title
plt.xlabel('Time (s)', fontsize=12)
plt.ylabel('Height (m)', fontsize=12)
plt.title("Trapezoid Method: Predictor-Corrector Structure", fontsize=14)
plt.legend(fontsize=10)
plt.grid(True, linestyle=':', alpha=0.7)

plt.tight_layout()
plt.show()
_images/86368566ab28dc13f7374c1baf8d42ed9d5ef7bde0a23d816d63afde84ca3713.png

5.3. Collocation Methods#

The numerical integration methods we’ve discussed before are inherently sequential: given an initial state, we make a guess as to where the system might go over a small time interval. The accuracy of that guess depends on the numerical integration procedure being used and the information available locally.

Collocation methods offer an alternative paradigm. Instead of solving the problem step by step, these methods aim to solve for the entire trajectory simultaneously by reformulating the differential equation into a system of algebraic equations. This approach involves a process of iterative refinements for the values of the discretized system, taking into account both local and global properties of the function.

Sequential methods like Euler’s or Runge-Kutta focus on evolving the system forward in time, starting from known initial conditions. They are simple to use for initial value problems but can accumulate errors over long time scales. Collocation methods, on the other hand, consider the whole time domain at once and tend to distribute errors more evenly across the domain.

This global view has some interesting implications. For one, collocation methods can handle both initial value and boundary value problems more naturally, which is particularly useful when you have path constraints throughout the entire integration interval. This property is especially valuable when solving continuous-time optimal control problems with such constraints: something which would otherwise be difficult to achieve with single-shooting methods.

However, solving for all time points simultaneously can be computationally intensive. It involves a tradeoff between overall accuracy, especially for long-time integrations or sensitive systems, but at the cost of increased computational complexity. Despite this, the ability to handle complex constraints and achieve more robust solutions often makes collocation methods the preferred choice for solving sophisticated COCPs in fields such as aerospace, robotics, and process control.

5.3.1. Quick Primer on Polynomials#

Collocation methods are based on polynomial approximation theory. Therefore, the first step in developing collocation-based optimal control techniques is to review the fundamentals of polynomial functions.

Polynomials are typically introduced through their standard form:

\[ p(t) = a_n t^n + a_{n-1} t^{n-1} + \cdots + a_1 t + a_0 \]

In this expression, the \(a_i\) are coefficients which linearly combine the powers of \(t\) to represent a function. The set of functions \(\{ 1, t, t^2, t^3, \ldots, t^n \}\) used in the standard polynomial representation is called the monomial basis.

In linear algebra, a basis is a set of vectors in a vector space such that any vector in the space can be uniquely represented as a linear combination of these basis vectors. In the same way, a polynomial basis is such that any function \( f(x) \) (within the function space) to be expressed as:

\[ f(x) = \sum_{k=0}^{\infty} c_k p_k(x), \]

where the coefficients \( c_k \) are generally determined by solving a system of equation.

Just as vectors can be represented in different coordinate systems (bases), functions can also be expressed using various polynomial bases. However, the ability to apply a change of basis does not imply that all types of polynomials are equivalent from a practical standpoint. In practice, our choice of polynomial basis is dictated by considerations of efficiency, accuracy, and stability when approximating a function.

For instance, despite the monomial basis being easy to understand and implement, it often performs poorly in practice due to numerical instability. This instability arises as its coefficients take on large values: an ill-conditioning problem. The following kinds of polynomial often remedy this issues.

5.3.1.1. Orthogonal Polynomials#

An orthogonal polynomial basis is a set of polynomials that are not only orthogonal to each other but also form a complete basis for a certain space of functions. This means that any function within that space can be represented as a linear combination of these polynomials.

More precisely, let \( \{ p_0(x), p_1(x), p_2(x), \dots \} \) be a sequence of polynomials where each \( p_n(x) \) is a polynomial of degree \( n \). We say that this set forms an orthogonal polynomial basis if any polynomial \( q(x) \) of degree \( n \) or less can be uniquely expressed as a linear combination of \( \{ p_0(x), p_1(x), \dots, p_n(x) \} \). Furthermore, the orthogonality property means that for any \( i \neq j \):

\[ \langle p_i, p_j \rangle = \int_a^b p_i(x) p_j(x) w(x) \, dx = 0. \]

for some weight function \( w(x) \) over a given interval of orthogonality \( [a, b] \).

The orthogonality property allows to simplify the computation of the coefficients involved in the polynomial representation of a function. At a high level, what happens is that when taking the inner product of \( f(x) \) with each basis polynomial, \( p_k(x) \) isolates the corresponding coefficient \( c_k \), which can be found to be:

\[ c_k = \frac{\langle f, p_k \rangle}{\langle p_k, p_k \rangle} = \frac{\int_a^b f(x) p_k(x) w(x) \, dx}{\int_a^b p_k(x)^2 w(x) \, dx}. \]

Here are some examples of the most common orthogonal polynomials used in practice.

5.3.1.1.1. Legendre Polynomials#

Legendre polynomials \( \{ P_n(x) \} \) are defined on the interval \([-1, 1]\) and satisfy the orthogonality condition:

\[\begin{split} \int_{-1}^{1} P_n(x) P_m(x) \, dx = \begin{cases} 0 & \text{if } n \neq m, \\ \frac{2}{2n + 1} & \text{if } n = m. \end{cases} \end{split}\]

They can be generated using the recurrence relation:

\[ (n+1) P_{n+1}(x) = (2n + 1) x P_n(x) - n P_{n-1}(x), \]

with initial conditions:

\[ P_0(x) = 1, \quad P_1(x) = x. \]

The first four Legendre polynomials resulting from this recurrence are the following:

Hide code cell source
import numpy as np
from IPython.display import display, Math

def legendre_polynomial(n, x):
    if n == 0:
        return np.poly1d([1])
    elif n == 1:
        return x
    else:
        p0 = np.poly1d([1])
        p1 = x
        for k in range(2, n + 1):
            p2 = ((2 * k - 1) * x * p1 - (k - 1) * p0) / k
            p0, p1 = p1, p2
        return p1

def legendre_coefficients(n):
    x = np.poly1d([1, 0])  # Define a poly1d object to represent x
    poly = legendre_polynomial(n, x)
    return poly

def poly_to_latex(poly):
    coeffs = poly.coefficients
    variable = poly.variable
    
    terms = []
    for i, coeff in enumerate(coeffs):
        power = len(coeffs) - i - 1
        if coeff == 0:
            continue
        coeff_str = f"{coeff:.2g}" if coeff not in {1, -1} or power == 0 else ("-" if coeff == -1 else "")
        if power == 0:
            term = f"{coeff_str}"
        elif power == 1:
            term = f"{coeff_str}{variable}"
        else:
            term = f"{coeff_str}{variable}^{power}"
        terms.append(term)
    
    latex_poly = " + ".join(terms).replace(" + -", " - ")
    return latex_poly

for n in range(4):
    poly = legendre_coefficients(n)
    display(Math(f"P_{n}(x) = {poly_to_latex(poly)}"))
\[\displaystyle P_0(x) = 1\]
\[\displaystyle P_1(x) = x\]
\[\displaystyle P_2(x) = 1.5x^2 - 0.5\]
\[\displaystyle P_3(x) = 2.5x^3 - 1.5x\]
5.3.1.1.2. Chebyshev Polynomials#

There are two types of Chebyshev polynomials: Chebyshev polynomials of the first kind, \( \{ T_n(x) \} \), and Chebyshev polynomials of the second kind, \( \{ U_n(x) \} \). We typically focus on the first kind. They are defined on the interval \([-1, 1]\) and satisfy the orthogonality condition:

\[\begin{split} \int_{-1}^{1} \frac{T_n(x) T_m(x)}{\sqrt{1 - x^2}} \, dx = \begin{cases} 0 & \text{if } n \neq m, \\ \frac{\pi}{2} & \text{if } n = m \neq 0, \\ \pi & \text{if } n = m = 0. \end{cases} \end{split}\]

The Chebyshev polynomials of the first kind can be generated using the recurrence relation:

\[ T_{n+1}(x) = 2x T_n(x) - T_{n-1}(x), \]

with initial conditions:

\[ T_0(x) = 1, \quad T_1(x) = x. \]

Remarkably, this recurrence relation also admits an explicit formula:

\[ T_n(x) = \cos(n \cos^{-1}(x)). \]

Let’s now implement it in Python:

Hide code cell source
def chebyshev_polynomial(n, x):
    if n == 0:
        return np.poly1d([1])
    elif n == 1:
        return x
    else:
        t0 = np.poly1d([1])
        t1 = x
        for _ in range(2, n + 1):
            t2 = 2 * x * t1 - t0
            t0, t1 = t1, t2
        return t1

def chebyshev_coefficients(n):
    x = np.poly1d([1, 0])  # Define a poly1d object to represent x
    poly = chebyshev_polynomial(n, x)
    return poly

for n in range(4):
    poly = chebyshev_coefficients(n)
    display(Math(f"T_{n}(x) = {poly_to_latex(poly)}"))
\[\displaystyle T_0(x) = 1\]
\[\displaystyle T_1(x) = x\]
\[\displaystyle T_2(x) = 2x^2 - 1\]
\[\displaystyle T_3(x) = 4x^3 - 3x\]
5.3.1.1.3. Hermite Polynomials#

Hermite polynomials \( \{ H_n(x) \} \) are defined on the entire real line and are orthogonal with respect to the weight function \( w(x) = e^{-x^2} \). They satisfy the orthogonality condition:

\[\begin{split} \int_{-\infty}^{\infty} H_n(x) H_m(x) e^{-x^2} \, dx = \begin{cases} 0 & \text{if } n \neq m, \\ 2^n n! \sqrt{\pi} & \text{if } n = m. \end{cases} \end{split}\]

Hermite polynomials can be generated using the recurrence relation:

\[ H_{n+1}(x) = 2x H_n(x) - 2n H_{n-1}(x), \]

with initial conditions:

\[ H_0(x) = 1, \quad H_1(x) = 2x. \]

The following code computes the coefficients of the first four Hermite polynomials:

Hide code cell source
def hermite_polynomial(n, x):
    if n == 0:
        return np.poly1d([1])
    elif n == 1:
        return 2 * x
    else:
        h0 = np.poly1d([1])
        h1 = 2 * x
        for k in range(2, n + 1):
            h2 = 2 * x * h1 - 2 * (k - 1) * h0
            h0, h1 = h1, h2
        return h1

def hermite_coefficients(n):
    x = np.poly1d([1, 0])  # Define a poly1d object to represent x
    poly = hermite_polynomial(n, x)
    return poly

for n in range(4):
    poly = hermite_coefficients(n)
    display(Math(f"H_{n}(x) = {poly_to_latex(poly)}"))
\[\displaystyle H_0(x) = 1\]
\[\displaystyle H_1(x) = 2x\]
\[\displaystyle H_2(x) = 4x^2 - 2\]
\[\displaystyle H_3(x) = 8x^3 - 12x\]

5.3.2. Collocation Conditions#

Consider a general ODE of the form:

\[ \dot{y}(t) = f(y(t), t), \quad y(t_0) = y_0, \]

where \( y(t) \in \mathbb{R}^n \) is the state vector, and \( f: \mathbb{R}^n \times \mathbb{R} \rightarrow \mathbb{R}^n \) is a known function. The goal is to approximate the solution \( y(t) \) over a given interval \([t_0, t_f]\). Collocation methods achieve this by:

  1. Choosing a basis to approximate \( y(t) \) using a finite sum of basis functions \( \phi_i(t) \):

    \[ y(t) \approx \sum_{i=0}^{N} c_i \phi_i(t), \]

    where \( \{c_i\} \) are the coefficients to be determined.

  2. Selecting collocation points \( t_1, t_2, \ldots, t_N \) within the interval \([t_0, t_f]\). These are typically chosen to be the roots of certain orthogonal polynomials, like Legendre or Chebyshev polynomials, or can be spread equally across the interval.

  3. Enforcing the ODE at the collocation points for each \( t_j \):

    \[ \dot{y}(t_j) = f(y(t_j), t_j). \]

    To implement this, we differentiate the approximate solution \( y(t) \) with respect to time:

    \[ \dot{y}(t) \approx \sum_{i=0}^{N} c_i \dot{\phi}_i(t). \]

    Substituting this into the ODE at the collocation points gives:

    \[ \sum_{i=0}^{N} c_i \dot{\phi}_i(t_j) = f\left(\sum_{i=0}^{N} c_i \phi_i(t_j), t_j\right), \quad j = 1, \ldots, N. \]

The collocation equations are formed by enforcing the ODE at all collocation points, leading to a system of nonlinear equations:

\[ \sum_{i=0}^{N} c_i \dot{\phi}_i(t_j) - f\left(\sum_{i=0}^{N} c_i \phi_i(t_j), t_j\right) = 0, \quad j = 1, \ldots, N. \]

Furthermore when solving an initial value problem (IVP), we also need to incorporate the initial condition \( y(t_0) = y_0 \) as an additional constraint:

\[ \sum_{i=0}^{N} c_i \phi_i(t_0) = y_0. \]

The collocation conditions and IVP condition are combined together to form a root-finding problem, which we can generically solve numerically using Newton’s method.

5.3.3. Common Numerical Integration Techniques as Collocation Methods#

Many common numerical integration techniques can be viewed as special cases of collocation methods. While the general collocation method we discussed earlier applies to the entire interval \([t_0, t_f]\), many numerical integration techniques can be viewed as collocation methods applied locally, step by step.

In practical numerical integration, we often divide the full interval \([t_0, t_f]\) into smaller subintervals or steps. In general, this allows us to user simpler basis functions thereby reducing computational complexity, and gives us more flexibility in dynamically ajusting the step size using local error estimates. When we apply collocation locally, we’re essentially using the collocation method to “step” from \(t_n\) to \(t_{n+1}\). As we did, earlier we still apply the following three steps:

  1. We choose a basis function to approximate \(y(t)\) over \([t_n, t_{n+1}]\).

  2. We select collocation points within this interval.

  3. We enforce the ODE at these points to determine the coefficients of our basis function.

We can make this idea clearer by re-deriving some of the numerical integration methods seen before using this perspective.

5.3.3.1. Explicit Euler Method#

For the Explicit Euler method, we use a linear basis function for each step:

\[ \phi(t) = 1 + c(t - t_n) \]

Note that we use \((t - t_n)\) rather than just \(t\) because we’re approximating the solution locally, relative to the start of each step. We then choose one collocation point at \(t_{n+1}\) where we have:

\[ y'(t_{n+1}) = c = f(y_n, t_n) \]

Our local approximation is:

\[ y(t) \approx y_n + c(t - t_n) \]

At \(t = t_{n+1}\), this gives:

\[ y_{n+1} = y_n + c(t_{n+1} - t_n) = y_n + hf(y_n, t_n) \]

where \(h = t_{n+1} - t_n\). This is the classic Euler update formula.

5.3.3.2. Implicit Euler Method#

The Implicit Euler method uses the same linear basis function:

\[ \phi(t) = 1 + c(t - t_n) \]

Again, we choose one collocation point at \(t_{n+1}\). The main difference is that we enforce the ODE using \(y_{n+1}\):

\[ y'(t_{n+1}) = c = f(y_{n+1}, t_{n+1}) \]

Our approximation remains:

\[ y(t) \approx y_n + c(t - t_n) \]

At \(t = t_{n+1}\), this leads to the implicit equation:

\[ y_{n+1} = y_n + hf(y_{n+1}, t_{n+1}) \]

5.3.3.3. Trapezoidal Method#

The Trapezoidal method uses a quadratic basis function:

\[ \phi(t) = 1 + c(t - t_n) + a(t - t_n)^2 \]

We use two collocation points: \(t_n\) and \(t_{n+1}\). Enforcing the ODE at these points gives:

  • At \(t_n\):

\[ y'(t_n) = c = f(y_n, t_n) \]
  • At \(t_{n+1}\):

\[ y'(t_{n+1}) = c + 2ah = f(y_n + ch + ah^2, t_{n+1}) \]

Our approximation is:

\[ y(t) \approx y_n + c(t - t_n) + a(t - t_n)^2 \]

At \(t = t_{n+1}\), this gives:

\[ y_{n+1} = y_n + ch + ah^2 \]

Solving the system of equations leads to the trapezoidal update:

\[ y_{n+1} = y_n + \frac{h}{2}[f(y_n, t_n) + f(y_{n+1}, t_{n+1})] \]

5.3.3.4. Runge-Kutta Methods#

Higher-order Runge-Kutta methods can also be interpreted as collocation methods. The RK4 method corresponds to a collocation method using a cubic polynomial basis:

\[ \phi(t) = 1 + c_1(t - t_n) + c_2(t - t_n)^2 + c_3(t - t_n)^3 \]

Here, we’re using a cubic polynomial to approximate the solution over each step, rather than the linear or quadratic approximations of the other methods above. For RK4, we use four collocation points:

  1. \(t_n\) (the start of the step)

  2. \(t_n + h/2\)

  3. \(t_n + h/2\)

  4. \(t_n + h\) (the end of the step)

These points are called the “Gauss-Lobatto” points, scaled to our interval \([t_n, t_n + h]\). The RK4 method enforces the ODE at these collocation points, leading to four stages:

\[\begin{split} \begin{aligned} k_1 &= hf(y_n, t_n) \\ k_2 &= hf(y_n + \frac{1}{2}k_1, t_n + \frac{h}{2}) \\ k_3 &= hf(y_n + \frac{1}{2}k_2, t_n + \frac{h}{2}) \\ k_4 &= hf(y_n + k_3, t_n + h) \end{aligned} \end{split}\]

The final update formula for RK4 can be derived by solving the system of equations resulting from enforcing the ODE at our collocation points:

\[ y_{n+1} = y_n + \frac{1}{6}(k_1 + 2k_2 + 2k_3 + k_4) \]

5.3.4. Example: Solving a Simple ODE by Collocation#

Consider a simple ODE:

\[ \frac{dy}{dt} = -y, \quad y(0) = 1, \quad t \in [0, 2] \]

The analytical solution is \(y(t) = e^{-t}\). We apply the collocation method with a monomial basis of order \(N\):

\[ \phi_i(t) = t^i, \quad i = 0, 1, \ldots, N \]

We select \(N\) equally spaced points \(\{t_1, \ldots, t_N\}\) in \([0, 2]\) as collocation points.

Hide code cell source
import numpy as np
from scipy.optimize import root
import matplotlib.pyplot as plt

def ode_function(y, t):
    """Define the ODE: dy/dt = -y"""
    return -y

def solve_ode_collocation(ode_func, t_span, y0, order):
    t0, tf = t_span
    n_points = order + 1  # number of collocation points
    t_points = np.linspace(t0, tf, n_points)
    
    def collocation_residuals(coeffs):
        residuals = []
        # Initial condition residual
        y_init = sum(c * t_points[0]**i for i, c in enumerate(coeffs))
        residuals.append(y_init - y0)
        # Collocation point residuals
        for t in t_points[1:]:  # Skip the first point as it's used for initial condition
            y = sum(c * t**i for i, c in enumerate(coeffs))
            dy_dt = sum(c * i * t**(i-1) for i, c in enumerate(coeffs) if i > 0)
            residuals.append(dy_dt - ode_func(y, t))
        return residuals

    # Initial guess for coefficients
    initial_coeffs = [y0] + [0] * order

    # Solve the system of equations
    solution = root(collocation_residuals, initial_coeffs)
    
    if not solution.success:
        raise ValueError("Failed to converge to a solution.")

    coeffs = solution.x

    # Generate solution
    t_fine = np.linspace(t0, tf, 100)
    y_solution = sum(c * t_fine**i for i, c in enumerate(coeffs))

    return t_fine, y_solution, t_points, coeffs

# Example usage
t_span = (0, 2)
y0 = 1
orders = [1, 2, 3, 4, 5]  # Different polynomial orders to try

plt.figure(figsize=(12, 8))

for order in orders:
    t, y, t_collocation, coeffs = solve_ode_collocation(ode_function, t_span, y0, order)
    
    # Calculate y values at collocation points
    y_collocation = sum(c * t_collocation**i for i, c in enumerate(coeffs))
    
    # Plot the results
    plt.plot(t, y, label=f'Order {order}')
    plt.scatter(t_collocation, y_collocation, s=50, zorder=5)

# Plot the analytical solution
t_analytical = np.linspace(t_span[0], t_span[1], 100)
y_analytical = y0 * np.exp(-t_analytical)
plt.plot(t_analytical, y_analytical, 'k--', label='Analytical')

plt.xlabel('t')
plt.ylabel('y')
plt.title('ODE Solutions: dy/dt = -y, y(0) = 1')
plt.legend()
plt.grid(True)
plt.show()

# Print error for each order
print("Maximum absolute errors:")
for order in orders:
    t, y, _, _ = solve_ode_collocation(ode_function, t_span, y0, order)
    y_true = y0 * np.exp(-t)
    max_error = np.max(np.abs(y - y_true))
    print(f"Order {order}: {max_error:.6f}")
_images/e8db85ef72932d4e8f3ea1a1c1cea0b8f24989863f529612870865441c521669.png
Maximum absolute errors:
Order 1: 0.300453
Order 2: 0.074113
Order 3: 0.015265
Order 4: 0.002518
Order 5: 0.000339

6. Trajectory Optimization in Continuous Time#

As studied earlier in the discrete-time setting, we consider three variants of the continuous-time optimal control problem (COCP) with path constraints and bounds:

Definition 6.1 (Mayer Problem)

\[\begin{split} \begin{aligned} \text{minimize} \quad & c(\mathbf{x}(t_f)) \\ \text{subject to} \quad & \dot{\mathbf{x}}(t) = \mathbf{f}(\mathbf{x}(t), \mathbf{u}(t)) \\ & \mathbf{g}(\mathbf{x}(t), \mathbf{u}(t)) \leq \mathbf{0} \\ & \mathbf{x}_{\text{min}} \leq \mathbf{x}(t) \leq \mathbf{x}_{\text{max}} \\ & \mathbf{u}_{\text{min}} \leq \mathbf{u}(t) \leq \mathbf{u}_{\text{max}} \\ \text{given} \quad & \mathbf{x}(t_0) = \mathbf{x}_0 \enspace . \end{aligned} \end{split}\]

Definition 6.2 (Lagrange Problem)

\[\begin{split} \begin{aligned} \text{minimize} \quad & \int_{t_0}^{t_f} c(\mathbf{x}(t), \mathbf{u}(t)) \, dt \\ \text{subject to} \quad & \dot{\mathbf{x}}(t) = \mathbf{f}(\mathbf{x}(t), \mathbf{u}(t)) \\ & \mathbf{g}(\mathbf{x}(t), \mathbf{u}(t)) \leq \mathbf{0} \\ & \mathbf{x}_{\text{min}} \leq \mathbf{x}(t) \leq \mathbf{x}_{\text{max}} \\ & \mathbf{u}_{\text{min}} \leq \mathbf{u}(t) \leq \mathbf{u}_{\text{max}} \\ \text{given} \quad & \mathbf{x}(t_0) = \mathbf{x}_0 \enspace . \end{aligned} \end{split}\]

Definition 6.3 (Bolza Problem)

\[\begin{split} \begin{aligned} \text{minimize} \quad & c(\mathbf{x}(t_f)) + \int_{t_0}^{t_f} c(\mathbf{x}(t), \mathbf{u}(t)) \, dt \\ \text{subject to} \quad & \dot{\mathbf{x}}(t) = \mathbf{f}(\mathbf{x}(t), \mathbf{u}(t)) \\ & \mathbf{g}(\mathbf{x}(t), \mathbf{u}(t)) \leq \mathbf{0} \\ & \mathbf{x}_{\text{min}} \leq \mathbf{x}(t) \leq \mathbf{x}_{\text{max}} \\ & \mathbf{u}_{\text{min}} \leq \mathbf{u}(t) \leq \mathbf{u}_{\text{max}} \\ \text{given} \quad & \mathbf{x}(t_0) = \mathbf{x}_0 \enspace . \end{aligned} \end{split}\]

In these formulations, the additional constraints are:

  • Path constraints: \(\mathbf{g}(\mathbf{x}(t), \mathbf{u}(t)) \leq \mathbf{0}\), which represent constraints that must be satisfied at all times along the trajectory.

  • State bounds: \(\mathbf{x}_{\text{min}} \leq \mathbf{x}(t) \leq \mathbf{x}_{\text{max}}\), which specify the lower and upper bounds on the state variables.

  • Control bounds: \(\mathbf{u}_{\text{min}} \leq \mathbf{u}(t) \leq \mathbf{u}_{\text{max}}\), which specify the lower and upper bounds on the control inputs.

Furthermore, we may also encounter variations of the above problems under the assumption that horizon is infinite. For example:

Definition 6.4 (Infinite-Horizon Trajectory Optimization)

\[\begin{align*} &\text{minimize} \quad \int_{t_0}^{\infty} e^{-\rho t} c(\mathbf{x}(t), \mathbf{u}(t)) \, dt \\ &\text{subject to} \quad \dot{\mathbf{x}}(t) = \mathbf{f}(\mathbf{x}(t), \mathbf{u}(t)) \\ &\phantom{\text{subject to}} \quad \mathbf{g}(\mathbf{x}(t), \mathbf{u}(t)) \leq \mathbf{0} \\ &\phantom{\text{subject to}} \quad \mathbf{x}_{\text{min}} \leq \mathbf{x}(t) \leq \mathbf{x}_{\text{max}} \\ &\phantom{\text{subject to}} \quad \mathbf{u}_{\text{min}} \leq \mathbf{u}(t) \leq \mathbf{u}_{\text{max}} \\ &\text{given} \quad \mathbf{x}(t_0) = \mathbf{x}_0 \enspace . \end{align*}\]

In this formulation, the term \(e^{-\rho t}\) is a discount factor that exponentially decreases the importance of future costs relative to the present. The parameter \( \rho > 0\) is the discount rate. A larger value of \( \rho \) places more emphasis on the immediate cost and diminishes the impact of future costs. In infinite-horizon problems, the integral of the cost function \( \int_{t_0}^{\infty} c(\mathbf{x}(t), \mathbf{u}(t)) \, dt \) could potentially diverge because the cost accumulates over an infinite time period. Introducing the exponential term \( e^{-\rho t} \) guarantees that the integral converges as long as \( c(\mathbf{x}(t), \mathbf{u}(t)) \) grows at a slower rate than \( e^{\rho t} \).

6.1. Transforming Lagrange and Bolza Problems to Mayer#

Recall that, according to the fundamental theorem of calculus, solving an initial value problem (IVP) or using a quadrature method to evaluate an integral are essentially two approaches to the same problem. Similarly, just as we can convert discrete-time Lagrange or Bolza optimal control problems to the Mayer form by keeping a running sum of the cost, we can extend this idea to continuous time. This is done by introducing an augmented state space that includes the accumulated cost up to a given time as part of the integral term. If our original system is represented by \(\mathbf{f}\), we construct an augmented system \(\tilde{\mathbf{f}}\) as follows:

\[\begin{split} \begin{aligned} \text{minimize} \quad & c(\mathbf{x}(t_f)) + z(t_f) \\ \text{subject to} \quad & \begin{bmatrix} \dot{\mathbf{x}}(t) \\ \dot{z}(t) \end{bmatrix} = \tilde{\mathbf{f}}(\mathbf{x}(t), z(t), \mathbf{u}(t)) \\ & \mathbf{g}(\mathbf{x}(t), \mathbf{u}(t)) \leq \mathbf{0} \\ & \mathbf{x}_{\text{min}} \leq \mathbf{x}(t) \leq \mathbf{x}_{\text{max}} \\ & \mathbf{u}_{\text{min}} \leq \mathbf{u}(t) \leq \mathbf{u}_{\text{max}} \\ \text{given} \quad & \mathbf{x}(t_0) = \mathbf{x}_0, \quad z(t_0) = 0 \enspace . \end{aligned} \end{split}\]

Here, the newly introduced state variable \(z(t)\) containts the accumulated integral cost over time. The augmented system dynamics, \(\tilde{\mathbf{f}}\), are defined as:

\[\begin{split} \tilde{\mathbf{f}}(\mathbf{x}, z, \mathbf{u}) = \begin{bmatrix} \mathbf{f}(\mathbf{x}, \mathbf{u}) \\ c(\mathbf{x}, \mathbf{u}) \end{bmatrix} \end{split}\]

The integral cost is now obtained via \(z(t_f)\), which is added with the terminal cost \(c(\mathbf{x}(t_f))\) to recover the value of the original objective in the Bolza problem.

6.2. Example Problems#

6.2.1. Inverted Pendulum#

The inverted pendulum is a classic problem in control theory and robotics that demonstrates the challenge of stabilizing a dynamic system that is inherently unstable. The objective is to keep a pendulum balanced in the upright position by applying a control force, typically at its base. This setup is analogous to balancing a broomstick on your finger: any deviation from the vertical position will cause the system to tip over unless you actively counteract it with appropriate control actions.

We typically assume that the pendulum is mounted on a cart or movable base, which can move horizontally. The system’s state is then characterized by four variables:

  1. Cart position: \( x(t) \) — the horizontal position of the base.

  2. Cart velocity: \( \dot{x}(t) \) — the speed of the cart.

  3. Pendulum angle: \( \theta(t) \) — the angle between the pendulum and the vertical upright position.

  4. Angular velocity: \( \dot{\theta}(t) \) — the rate at which the pendulum’s angle is changing.

This setup is more complex because the controller must deal with interactions between two different types of motion: linear (the cart) and rotational (the pendulum). This system is said to be “underactuated” because the number of control inputs (one) is less than the number of state variables (four). This makes the problem more challenging and interesting from a control perspective.

We can simplify the problem by assuming that the base of the pendulum is fixed. This is akin to having the bottom of the stick attached to a fixed pivot on a table. You can’t move the base anymore; you can only apply small nudges at the pivot point to keep the stick balanced upright. In this case, you’re only focusing on adjusting the stick’s tilt without worrying about moving the base. This reduces the problem to stabilizing the pendulum’s upright orientation using only the rotational dynamics. The system’s state can now be described by just two variables:

  1. Pendulum angle: \( \theta(t) \) — the angle of the pendulum from the upright vertical position.

  2. Angular velocity: \( \dot{\theta}(t) \) — the rate at which the pendulum’s angle is changing.

The evolution of these two varibles is governed by the following ordinary differential equation:

(6.1)#\[\begin{equation} \begin{bmatrix} \dot{\theta}(t) \\ \ddot{\theta}(t) \end{bmatrix} = \begin{bmatrix} \dot{\theta}(t) \\ \frac{mgl}{J_t} \sin{\theta(t)} - \frac{\gamma}{J_t} \dot{\theta}(t) + \frac{l}{J_t} u(t) \cos{\theta(t)} \end{bmatrix}, \quad y(t) = \theta(t) \end{equation}\]

where:

  • \(m\) is the mass of the pendulum

  • \(g\) is the acceleration due to gravity

  • \(l\) is the length of the pendulum

  • \(\gamma\) is the coefficient of rotational friction

  • \(J_t = J + ml^2\) is the total moment of inertia, with \(J\) being the pendulum’s moment of inertia about its center of mass

  • \(u(t)\) is the control force applied at the base

  • \(y(t) = \theta(t)\) is the measured output (the pendulum’s angle)

We expect that when no control is applied to the system, the rod should be falling down when started from the upright position.

Hide code cell source
import numpy as np
import matplotlib.pyplot as plt
from scipy.integrate import odeint
from IPython.display import HTML, display
from matplotlib.animation import FuncAnimation

# System parameters
m = 1.0  # mass of the pendulum (kg)
g = 9.81  # acceleration due to gravity (m/s^2)
l = 1.0  # length of the pendulum (m)
gamma = 0.1  # coefficient of rotational friction
J = 1/3 * m * l**2  # moment of inertia of a rod about its center of mass
J_t = J + m * l**2  # total moment of inertia

# Define the ODE for the inverted pendulum
def pendulum_ode(state, t):
    theta, omega = state
    dtheta = omega
    domega = (m*g*l/J_t) * np.sin(theta) - (gamma/J_t) * omega
    return [dtheta, domega]

# Initial conditions: slightly off vertical position
theta0 = 0.1  # initial angle (radians)
omega0 = 0  # initial angular velocity (rad/s)
y0 = [theta0, omega0]

# Time array for integration
t = np.linspace(0, 10, 500)  # Reduced number of points

# Solve ODE
solution = odeint(pendulum_ode, y0, t)

# Extract theta and omega from the solution
theta = solution[:, 0]
omega = solution[:, 1]

# Create two separate plots
fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(10, 10))

# Plot for angle
ax1.plot(t, theta)
ax1.set_xlabel('Time (s)')
ax1.set_ylabel('Angle (rad)')
ax1.set_title('Pendulum Angle over Time')
ax1.grid(True)

# Plot for angular velocity
ax2.plot(t, omega)
ax2.set_xlabel('Time (s)')
ax2.set_ylabel('Angular velocity (rad/s)')
ax2.set_title('Pendulum Angular Velocity over Time')
ax2.grid(True)

plt.tight_layout()
plt.show()

# Function to create animation frames
def get_pendulum_position(theta):
    x = l * np.sin(theta)
    y = l * np.cos(theta)
    return x, y

# Create animation
fig, ax = plt.subplots(figsize=(8, 8))
line, = ax.plot([], [], 'o-', lw=2)
time_text = ax.text(0.02, 0.95, '', transform=ax.transAxes)

def init():
    ax.set_xlim(-1.2*l, 1.2*l)
    ax.set_ylim(-1.2*l, 1.2*l)  # Adjusted to show full range of motion
    ax.set_aspect('equal', adjustable='box')
    return line, time_text

def animate(i):
    x, y = get_pendulum_position(theta[i])
    line.set_data([0, x], [0, y])
    time_text.set_text(f'Time: {t[i]:.2f} s')
    return line, time_text

anim = FuncAnimation(fig, animate, init_func=init, frames=len(t), interval=40, blit=True)
plt.title('Inverted Pendulum Animation')
ax.grid(True)

# Convert animation to JavaScript
js_anim = anim.to_jshtml()

# Close the figure to prevent it from being displayed
plt.close(fig)

# Display only the JavaScript animation
display(HTML(js_anim))
_images/8e9b683de8f201d7a158cb8bcf4a1f449df582650e19d9b82fde6b1da181ff14.png

6.2.1.1. Looking Under the Hood: Pendulum in the Gym Environment#

Gym is a widely used abstraction layer for defining discrete-time reinforcement learning problems. In reinforcement learning research, there’s often a desire to develop general-purpose algorithms that are problem-agnostic. This research mindset leads us to voluntarily avoid considering the implementation details of a given environment. While this approach is understandable from a research perspective, it may not be optimal from a pragmatic, solution-driven standpoint where we care about solving specific problems efficiently. If we genuinely wanted to solve this problem without prior knowledge, why not look under the hood and embrace its nature as a trajectory optimization problem?

Let’s examine the code and reverse-engineer the original continuous-time problem hidden behind the abstraction layer. Although the pendulum problem may have limited practical relevance as a real-world application, it serves as an excellent example for our analysis. In the current version of Pendulum, we find that the Gym implementation uses a simplified model. Like our implementation, it assumes a fixed base and doesn’t model cart movement. The state is also represented by the pendulum angle and angular velocity. However, the equations of motion implemented in the Gym environment are different and correspond to the following ODE:

\[\begin{align*} \dot{\theta} &= \theta_{dot} \\ \dot{\theta}_{dot} &= \frac{3g}{2l} \sin(\theta) + \frac{3}{ml^2} u \end{align*}\]

Compared to our simplified model, the Gym implementation makes the following additional assumptions:

  1. It omits the term \(\frac{\gamma}{J_t} \dot{\theta}(t)\), which represents damping or air resistance. This means that it assumes an idealized pendulum that doesn’t naturally slow down over time.

  2. It uses \(ml^2\) instead of \(J_t = J + ml^2\), which assumes that all mass is concentrated at the pendulum’s end (like a point mass on a massless rod), rather than accounting for mass distribution along the pendulum.

  3. The control input \(u\) is applied directly, without a \(\cos \theta(t)\) term, which means that the applied torque has the same effect regardless of the pendulum’s position, rather than varying with angle. For example, imagine trying to push a door open. When the door is almost closed (pendulum near vertical), a small push perpendicular to the door (analogous to our control input) can easily start it moving. However, when the door is already wide open (pendulum horizontal), the same push has little effect on the door’s angle. In a more detailed model, this would be captured by the \(\cos \theta(t)\) term, which is maximum when the pendulum is vertical (\(\cos 0° = 1\)) and zero when horizontal (\(\cos 90° = 0\)).

The goal remains to stabilize the rod upright, but the way in which this encoded is through the following instantenous cost function:

\[\begin{align*} c(\theta, \dot{\theta}, u) &= (\text{normalize}(\theta))^2 + 0.1\dot{\theta}^2 + 0.001u^2\\ \text{normalize}(\theta) &= ((\theta + \pi) \bmod 2\pi) - \pi \end{align*}\]

This cost function penalizes deviations from the upright position (first term), discouraging rapid motion (second term), and limiting control effort (third term). The relative weights has been manually chosen to balance the primary goal of upright stabilization with the secondary aims of smooth motion and energy efficiency. The normalization ensures that the angle is always in the range \([-\pi, \pi]\) so that the pendulum positions (e.g., \(0\) and \(2\pi\)) are treated identically, which could otherwise confuse learning algorithms.

Studying the code further, we find that it imposes bound constraints on both the control input and the angular velocity through clipping operations:

\[\begin{align*} u &= \max(\min(u, u_{max}), -u_{max}) \\ \dot{\theta} &= \max(\min(\dot{\theta}, \dot{\theta}_{max}), -\dot{\theta}_{max}) \end{align*}\]

Where \(u_{max} = 2.0\) and \(\dot{\theta}_{max} = 8.0\). Finally, when inspecting the step function, we find that the dynamics are discretized using forward Euler under a fixed step size of \(h=0.0.5\). Overall, the discrete-time trajectory optimization problem implemented in Gym is the following:

\[\begin{align*} \min_{u_k} \quad & J = \sum_{k=0}^{N-1} c(\theta_k, \dot{\theta}_k, u_k) \\ \text{subject to:} \quad & \theta_{k+1} = \theta_k + \dot{\theta}_k \cdot h \\ & \dot{\theta}_{k+1} = \dot{\theta}_k + \left(\frac{3g}{2l}\sin(\theta_k) + \frac{3}{ml^2}u_k\right) \cdot h \\ & -u_{\max} \leq u_k \leq u_{\max} \\ & -\dot{\theta}_{\max} \leq \dot{\theta}_k \leq \dot{\theta}_{\max}, \quad k = 0, 1, ..., N-1 \\ \text{given:} \quad & \theta_0 = \theta_{\text{initial}}, \quad \dot{\theta}_0 = \dot{\theta}_{\text{initial}}, \quad N = 200 \end{align*}\]

with \(g = 10.0\), \(l = 1.0\), \(m = 1.0\), \(u_{max} = 2.0\), and \(\dot{\theta}_{max} = 8.0\). This discrete-time problem corresponds to the following continuous-time optimal control problem:

\[\begin{align*} \min_{u(t)} \quad & J = \int_{0}^{T} c(\theta(t), \dot{\theta}(t), u(t)) dt \\ \text{subject to:} \quad & \dot{\theta}(t) = \dot{\theta}(t) \\ & \ddot{\theta}(t) = \frac{3g}{2l}\sin(\theta(t)) + \frac{3}{ml^2}u(t) \\ & -u_{\max} \leq u(t) \leq u_{\max} \\ & -\dot{\theta}_{\max} \leq \dot{\theta}(t) \leq \dot{\theta}_{\max} \\ \text{given:} \quad & \theta(0) = \theta_0, \quad \dot{\theta}(0) = \dot{\theta}_0, \quad T = 10 \text{ seconds} \end{align*}\]

6.2.2. Heat Exchanger#

Heat Exchanger

We are considering a system where fluid flows through a tube, and the goal is to control the temperature of the fluid by adjusting the temperature of the tube’s wall over time. The wall temperature, denoted as \( T_w(t) \), can be changed as a function of time, but it remains the same along the length of the tube. On the other hand, the temperature of the fluid inside the tube, \( T(z, t) \), depends both on its position along the tube \( z \) and on time \( t \). It evolves according to the following partial differential equation:

\[ \frac{\partial T}{\partial t} = -v \frac{\partial T}{\partial z} + \frac{h}{\rho C_p} (T_w(t) - T) \]

where we have:

  • \( v \): the average speed of the fluid moving through the tube,

  • \( h \): how easily heat transfers from the wall to the fluid,

  • \( \rho \) and \( C_p \): the fluid’s density and heat capacity.

This equation describes how the fluid’s temperature changes as it moves along the tube and interacts with the tube’s wall temperature. The fluid enters the tube with an initial temperature \( T_0 \) at the inlet (where \( z = 0 \)). Our objective is to adjust the wall temperature \( T_w(t) \) so that by a specific final time \( t_f \), the fluid’s temperature reaches a desired distribution \( T_s(z) \) along the length of the tube. The relationship for \( T_s(z) \) under steady-state conditions (ie. when changes over time are no longer considered), is given by:

\[ \frac{d T_s}{d z} = \frac{h}{v \rho C_p}[\theta - T_s] \]

where \( \theta \) is a constant temperature we want to maintain at the wall. The objective is to control the wall temperature \( T_w(t) \) so that by the end of the time interval \( t_f \), the fluid temperature \( T(z, t_f) \) is as close as possible to the desired distribution \( T_s(z) \). This can be formalized by minimizing the following quantity:

\[ I = \int_0^L \left[T(z, t_f) - T_s(z)\right]^2 dz \]

where \( L \) is the length of the tube. Additionally, we require that the wall temperature cannot exceed a maximum allowable value \( T_{\max} \):

\[ T_w(t) \leq T_{\max} \]

6.2.3. Nuclear Reactor#

Nuclear Reactor Diagram

In a nuclear reactor, neutrons interact with fissile nuclei, causing nuclear fission. This process produces more neutrons and smaller fissile nuclei called precursors. The precursors subsequently absorb more neutrons, generating “delayed” neutrons. The kinetic energy of these products is converted into thermal energy through collisions with neighboring atoms. The reactor’s power output is determined by the concentration of neutrons available for nuclear fission.

The reaction kinetics can be modeled using a system of ordinary differential equations:

\[\begin{align*} \dot{x}(t) &= \frac{r(t)x(t) - \alpha x^2(t) - \beta x(t)}{\tau} + \mu y(t), & x(0) &= x_0 \\ \dot{y}(t) &= \frac{\beta x(t)}{\tau} - \mu y(t), & y(0) &= y_0 \end{align*}\]

where:

  • \(x(t)\): concentration of neutrons at time \(t\)

  • \(y(t)\): concentration of precursors at time \(t\)

  • \(t\): time

  • \(r(t) = r[u(t)]\): degree of change in neutron multiplication at time \(t\) as a function of control rod displacement \(u(t)\)

  • \(\alpha\): reactivity coefficient

  • \(\beta\): fraction of delayed neutrons

  • \(\mu\): decay constant for precursors

  • \(\tau\): average time taken by a neutron to produce a neutron or precursor

The power output can be adjusted based on demand by inserting or retracting a neutron-absorbing control rod. Inserting the control rod absorbs neutrons, reducing the heat flux and power output, while retracting the rod has the opposite effect.

The objective is to change the neutron concentration \(x(t)\) from an initial value \(x_0\) to a stable value \(x_\mathrm{f}\) at time \(t_\mathrm{f}\) while minimizing the displacement of the control rod. This can be formulated as an optimal control problem, where the goal is to find the control function \(u(t)\) that minimizes the objective functional:

\[\begin{equation*} I = \int_0^{t_\mathrm{f}} u^2(t) \, \mathrm{d}t \end{equation*}\]

subject to the final conditions:

\[\begin{align*} x(t_\mathrm{f}) &= x_\mathrm{f} \\ \dot{x}(t_\mathrm{f}) &= 0 \end{align*}\]

and the constraint \(|u(t)| \leq u_\mathrm{max}\)

6.2.4. Chemotherapy#

Chemotherapy uses drugs to kill cancer cells. However, these drugs can also have toxic effects on healthy cells in the body. To optimize the effectiveness of chemotherapy while minimizing its side effects, we can formulate an optimal control problem.

The drug concentration \(y_1(t)\) and the number of immune cells \(y_2(t)\), healthy cells \(y_3(t)\), and cancer cells \(y_4(t)\) in an organ at any time \(t\) during chemotherapy can be modeled using a system of ordinary differential equations:

\[\begin{align*} \dot{y}_1(t) &= u(t) - \gamma_6 y_1(t) \\ \dot{y}_2(t) &= \dot{y}_{2,\text{in}} + r_2 \frac{y_2(t) y_4(t)}{\beta_2 + y_4(t)} - \gamma_3 y_2(t) y_4(t) - \gamma_4 y_2(t) - \alpha_2 y_2(t) \left(1 - e^{-y_1(t) \lambda_2}\right) \\ \dot{y}_3(t) &= r_3 y_3(t) \left(1 - \beta_3 y_3(t)\right) - \gamma_5 y_3(t) y_4(t) - \alpha_3 y_3(t) \left(1 - e^{-y_1(t) \lambda_3}\right) \\ \dot{y}_4(t) &= r_1 y_4(t) \left(1 - \beta_1 y_4(t)\right) - \gamma_1 y_3(t) y_4(t) - \gamma_2 y_2(t) y_4(t) - \alpha_1 y_4(t) \left(1 - e^{-y_1(t) \lambda_1}\right) \end{align*}\]

where:

  • \(y_1(t)\): drug concentration in the organ at time \(t\)

  • \(y_2(t)\): number of immune cells in the organ at time \(t\)

  • \(y_3(t)\): number of healthy cells in the organ at time \(t\)

  • \(y_4(t)\): number of cancer cells in the organ at time \(t\)

  • \(\dot{y}_{2,\text{in}}\): constant rate of immune cells entering the organ to fight cancer cells

  • \(u(t)\): rate of drug injection into the organ at time \(t\)

  • \(r_i, \beta_i\): constants in the growth terms

  • \(\alpha_i, \lambda_i\): constants in the decay terms due to the action of the drug

  • \(\gamma_i\): constants in the remaining decay terms

The objective is to minimize the number of cancer cells \(y_4(t)\) in a specified time \(t_\mathrm{f}\) while using the minimum amount of drug to reduce its toxic effects. This can be formulated as an optimal control problem, where the goal is to find the control function \(u(t)\) that minimizes the objective functional:

\[\begin{equation*} I = y_4(t_\mathrm{f}) + \int_0^{t_\mathrm{f}} u(t) \, \mathrm{d}t \end{equation*}\]

subject to the system dynamics, initial conditions, and the constraint \(u(t) \geq 0\).

Additional constraints may include:

  • Maintaining a minimum number of healthy cells during treatment:

    \[\begin{equation*} y_3(t) \geq y_{3,\min} \end{equation*}\]
  • Imposing an upper limit on the drug dosage:

    \[\begin{equation*} u(t) \leq u_\max \end{equation*}\]

6.2.5. Government Corruption#

In this model from Feichtinger and Wirl (1994), we aim to understand the incentives for politicians to engage in corrupt activities or to combat corruption. The model considers a politician’s popularity as a dynamic process that is influenced by the public’s memory of recent and past corruption. The objective is to find conditions under which self-interested politicians would choose to be honest or dishonest.

The model introduces the following notation:

  • \(C(t)\): accumulated awareness (knowledge) of past corruption at time \(t\)

  • \(u(t)\): extent of corruption (politician’s control variable) at time \(t\)

  • \(\delta\): rate of forgetting past corruption

  • \(P(t)\): politician’s popularity at time \(t\)

  • \(g(P)\): growth function of popularity; \(g''(P) < 0\)

  • \(f(C)\): function measuring the loss of popularity caused by \(C\); \(f'(C) > 0\), \(f''(C) \geq 0\)

  • \(U_1(P)\): benefits associated with being popular; \(U_1'(P) > 0\), \(U_1''(P) \leq 0\)

  • \(U_2(u)\): benefits resulting from bribery and fraud; \(U_2'(u) > 0\), \(U_2''(u) < 0\)

  • \(r\): discount rate

The dynamics of the public’s memory of recent and past corruption \(C(t)\) are modeled as:

\[\begin{align*} \dot{C}(t) &= u(t) - \delta C(t), \quad C(0) = C_0 \end{align*}\]

The evolution of the politician’s popularity \(P(t)\) is governed by:

\[\begin{align*} \dot{P}(t) &= g(P(t)) - f(C(t)), \quad P(0) = P_0 \end{align*}\]

The politician’s objective is to maximize the following objective:

\[\begin{equation*} \int_0^{\infty} e^{-rt} [U_1(P(t)) + U_2(u(t))] \, \mathrm{d}t \end{equation*}\]

subject to the dynamics of corruption awareness and popularity.

The optimal control problem can be formulated as follows:

\[\begin{align*} \max_{u(\cdot)} \quad & \int_0^{\infty} e^{-rt} [U_1(P(t)) + U_2(u(t))] \, \mathrm{d}t \\ \text{s.t.} \quad & \dot{C}(t) = u(t) - \delta C(t), \quad C(0) = C_0 \\ & \dot{P}(t) = g(P(t)) - f(C(t)), \quad P(0) = P_0 \end{align*}\]

The state variables are the accumulated awareness of past corruption \(C(t)\) and the politician’s popularity \(P(t)\). The control variable is the extent of corruption \(u(t)\). The objective functional represents the discounted stream of benefits coming from being honest (popularity) and from being dishonest (corruption).

6.3. Direct Transcription Methods#

When transitioning from discrete-time to continuous-time optimal control, several new computational challenges arise:

  1. The optimization variables are now continuous functions, not just discrete sequences of values stored in an array. We seek \(x(t)\) and \(u(t)\), which are continuous functions of time.

  2. Evaluating a candidate pair \(x(t)\) and \(u(t)\) requires integration. This is necessary both for computing the integral term in the objective of Lagrange or Bolza problems and for satisfying the dynamics expressed as constraints.

These challenges can be viewed as problems of representation and integration. Representation is addressed through function approximation, while integration is handled using numerical integration methods. In fields like aeronautics, aerospace, and chemical engineering, control or state functions are often represented using polynomials. Instead of searching directly over the space of all possible functions, we search over the space of parameters defining these polynomials. This approach is similar to deep learning, where we parameterize complex functions as compositions of simpler nonlinear functions and adjust their weights through gradient descent, rather than searching the function space directly. Therefore, the state and control parameterization technique developed in this chapter, using polynomials, can be extended to other function approximation methods, including neural networks.

Evaluating the integral is essentially a problem of numerical integration, which naturally pairs with the numerical integration of the dynamics. Together, these elements form the foundation of many methods collectively known as direct transcription methods. The key idea is to transform the original continuous-time optimal control problem—an infinite-dimensional optimization problem—into a finite-dimensional approximation that can be solved as a standard nonlinear programming (NLP) problem, similar to those we have encountered in discrete-time optimal control.

6.3.1. Direct Single Shooting#

Single shooting under a direct transcription approach amounts to expressing the equality constraints of the original problem about the ODE dynamics by a time-discretized counterpart unrolled in time. This method essentially “backpropagates” through the numerical integration code that implements the system.

6.3.1.1. Mayer Problem#

Consider a problem in Mayer form:

\[\begin{split} \begin{aligned} \text{minimize} \quad & c(\mathbf{x}(t_f)) \\ \text{subject to} \quad & \dot{\mathbf{x}}(t) = \mathbf{f}(\mathbf{x}(t), \mathbf{u}(t)), \quad t \in [t_0, t_f] \\ & \mathbf{u}_{\text{min}} \leq \mathbf{u}(t) \leq \mathbf{u}_{\text{max}}, \quad t \in [t_0, t_f] \\ \text{given} \quad & \mathbf{x}(t_0) = \mathbf{x}_0 \enspace . \end{aligned} \end{split}\]

Our first step is to pick a control parameterization. Traditionally this would be polynomials, but given the ease of coding these methods with automatic differentiation, it might as well be a neural network. Therefore, let’s write \(\mathbf{u}(t; \boldsymbol{\theta})\) to denote the parameterized controls. Second, let’s eliminate the approximate dynamics as constraints through a process of “simulation” via the chosen time discretization scheme. For RK4, we would for example obtain:

\[\begin{split} \begin{aligned} \text{minimize}_{\boldsymbol{\theta}} \quad & c(\Phi(\boldsymbol{\theta}; \mathbf{x}_0)) \\ \text{subject to} \quad & \mathbf{u}_{\text{min}} \leq \mathbf{u}(t_i; \boldsymbol{\theta}) \leq \mathbf{u}_{\text{max}}, \quad i = 0, \ldots, N-1 \\ \text{where} \quad & \Phi = \Psi_N \circ \Psi_{N-1} \circ \cdots \circ \Psi_1 \\ \Psi_i(\mathbf{x}) &= \mathbf{x} + \frac{\Delta t}{6}(k_1 + 2k_2 + 2k_3 + k_4) \\ k_1 &= \mathbf{f}(\mathbf{x}, \mathbf{u}(t_i; \boldsymbol{\theta})) \\ k_2 &= \mathbf{f}(\mathbf{x} + \frac{\Delta t}{2}k_1, \mathbf{u}(t_i + \frac{\Delta t}{2}; \boldsymbol{\theta})) \\ k_3 &= \mathbf{f}(\mathbf{x} + \frac{\Delta t}{2}k_2, \mathbf{u}(t_i + \frac{\Delta t}{2}; \boldsymbol{\theta})) \\ k_4 &= \mathbf{f}(\mathbf{x} + \Delta t k_3, \mathbf{u}(t_i + \Delta t; \boldsymbol{\theta})) \end{aligned} \end{split}\]

Here, \(\Phi(\boldsymbol{\theta}; \mathbf{x}_0)\) represents the final state \(\mathbf{x}(t_f)\) obtained by integrating the system dynamics from the initial state \(\mathbf{x}_0\) using the parameterized control \(\mathbf{u}(t; \boldsymbol{\theta})\). The functions \(\Psi_i\) represent the numerical integration steps (in this case, using the fourth-order Runge-Kutta method) for each time interval.

This reformulation transforms the infinite-dimensional optimal control problem into a finite-dimensional nonlinear programming (NLP) problem. The decision variables are now the parameters \(\boldsymbol{\theta}\) of the control function, subject to bound constraints at each discretization point. The dynamic constraints are implicitly satisfied through the integration process embedded in the objective function. However, ensuring that bound constraints on the controls are satisfied by a choice of parameters is more challenging than in the “tabular” case, i.e., when we are not using a control parameterization. This makes it more difficult to directly project or clip the parameters to satisfy the control bounds through a simple projected gradient descent step.

6.3.1.2. Bolza Problem#

Let’s now address the numerical challenge that comes with the evaluation of the integral term in a Lagrange or Bolza-type problem:

\[\begin{split} \begin{aligned} \text{minimize} \quad & c(\mathbf{x}(t_f)) + \int_{t_0}^{t_f} c(\mathbf{x}(t), \mathbf{u}(t)) \, dt \\ \text{subject to} \quad & \dot{\mathbf{x}}(t) = \mathbf{f}(\mathbf{x}(t), \mathbf{u}(t)) \\ & \mathbf{g}(\mathbf{x}(t), \mathbf{u}(t)) \leq \mathbf{0} \\ & \mathbf{u}_{\text{min}} \leq \mathbf{u}(t) \leq \mathbf{u}_{\text{max}} \\ \text{given} \quad & \mathbf{x}(t_0) = \mathbf{x}_0 \enspace . \end{aligned} \end{split}\]

We can first start with the usual trick by which we transform the above Bolza problem into Mayer form by creating an augmented system and solving the following COCP:

\[\begin{split} \begin{aligned} \text{minimize} \quad & c(\mathbf{x}(t_f)) + z(t_f) \\ \text{subject to} \quad & \begin{bmatrix} \dot{\mathbf{x}}(t) \\ \dot{z}(t) \end{bmatrix} = \tilde{\mathbf{f}}(\mathbf{x}(t), z(t), \mathbf{u}(t)) \\ & \mathbf{g}(\mathbf{x}(t), \mathbf{u}(t)) \leq \mathbf{0} \\ & \mathbf{u}_{\text{min}} \leq \mathbf{u}(t) \leq \mathbf{u}_{\text{max}} \\ \text{where} \quad & \tilde{\mathbf{f}}(\mathbf{x}, z, \mathbf{u}) = \begin{bmatrix} \mathbf{f}(\mathbf{x}, \mathbf{u}) \\ c(\mathbf{x}, \mathbf{u}) \end{bmatrix}\\ \text{given} \quad & \mathbf{x}(t_0) = \mathbf{x}_0, \quad z(t_0) = 0 \enspace . \end{aligned} \end{split}\]

We can then proceed to transcribe this problem via single-shooting just like we did above for Mayer problems:

\[\begin{split} \begin{aligned} \text{minimize}_{\boldsymbol{\theta}} \quad & c(\Phi(\boldsymbol{\theta}; \mathbf{x}_0)) + z_N \\ \text{subject to} \quad & \mathbf{u}_{\text{min}} \leq \mathbf{u}(t_i; \boldsymbol{\theta}) \leq \mathbf{u}_{\text{max}}, \quad i = 0, \ldots, N-1 \\ \text{where} \quad & \tilde{\Phi} = \tilde{\Psi}_N \circ \tilde{\Psi}_{N-1} \circ \cdots \circ \tilde{\Psi}_1 \\ \quad & \tilde{\Psi}_i \left( \begin{bmatrix} \mathbf{x}_{i-1} \\ z_{i-1} \end{bmatrix}, \mathbf{u}_{i-1} \right) = \begin{bmatrix} \mathbf{x}_i \\ z_i \end{bmatrix}\quad \text{for} \quad i = 1, \ldots, N \enspace , \end{aligned} \end{split}\]

where the \(\tilde \Psi\) denote a step of the underlying numerical integration method (eg. Euler, RK4, etc.) over the augmented system \(\tilde{\mathbf{f}}\).

6.3.1.3. Decoupled Integration#

The transformation of the Bolza problem into a Mayer problem has a side effect that it couples the numerical integration scheme used for the dynamics and the one for the objective itself. It certainly simplifies our implementation, but at the cost of less flexibility in the approximation of the problem as we will see below. We could for example directly apply a numerical quadrature method (say gaussian or simpson quadrature) and RKH4, resulting in the following transcribed (single shooted) problem:

\[\begin{split} \begin{aligned} \text{minimize}{\boldsymbol{\theta}} \quad & c(\Phi_N(\boldsymbol{\theta}; \mathbf{x}_0)) + Q(\boldsymbol{\theta}) \\ \text{subject to} \quad & \mathbf{u}_{\text{min}} \leq \mathbf{u}(t_i; \boldsymbol{\theta}) \leq \mathbf{u}_{\text{max}}, \quad i = 0, \ldots, N-1 \\ \text{where} \quad & \Phi_i = \Psi_i \circ \Psi_{i-1} \circ \cdots \circ \Psi_1 \\ \quad & Q(\boldsymbol{\theta}) = \sum_{i=0}^{N} w_i c(\Phi_i(\boldsymbol{\theta}; \mathbf{x}_0), \mathbf{u}(t_i; \boldsymbol{\theta})) \end{aligned} \end{split}\]

Here, \(\Phi_i\) represents the integration of the system dynamics using RK4 as before, while \(Q(\boldsymbol{\theta})\) represents the approximation of the integral cost using a chosen quadrature method. The weights \(w_i\) and the number of points \(N\) depend on the chosen quadrature method. For example:

  1. For the trapezoidal rule, we have N+1 equally spaced points, and the weights are:

\[ w_0 = w_N = \frac{\Delta t}{2}, \quad w_i = \Delta t \text{ for } i = 1, \ldots, N-1 \]

where \(\Delta t = (t_f - t_0) / N\) is the time step.

  1. For Simpson’s rule (assuming \(N\) is even), we have N+1 equally spaced points, and the weights are:

\[ w_0 = w_N = \frac{\Delta t}{3}, \quad w_i = \frac{4\Delta t}{3} \text{ for odd } i, \quad w_i = \frac{2\Delta t}{3} \text{ for even } i \neq 0, N \]
  1. For 2-point Gaussian quadrature, we have \(2N\) points (two per subinterval), and the weights are:

\[ w_{2i-1} = w_{2i} = \frac{t_f - t_0}{2N} \text{ for } i = 1, \ldots, N \]

The evaluation points \(t_i\) are not equally spaced in this case, but are determined by the Gaussian quadrature formula within each subinterval.

In all these cases, the states at the quadrature points are obtained from the integration of the system dynamics using the chosen method (e.g., RK4), represented by \(\Phi_i(\boldsymbol{\theta}; \mathbf{x}_0)\). However, it’s important to note that this approach can lead to computational inefficiencies if we’re not careful. If the quadrature points don’t align with the grid points used for integrating the system dynamics, we may need to perform additional interpolations or integrations to obtain the state values at the quadrature points. This mismatch can significantly increase the computational cost, especially for high-order quadrature methods or fine time discretizations.

For instance, when using Gaussian quadrature, the evaluation points are generally not equally spaced and won’t coincide with the uniform time grid typically used for RK4 integration. This means we might need to integrate the dynamics to intermediate points between our main time steps, thereby increasing the number of integration steps required. Similarly, for Simpson’s rule or other higher-order methods, we might need state values at points where we haven’t directly computed them during our primary integration process.

To mitigate this issue, one could align the integration grid with the quadrature points, but this might compromise the accuracy of the dynamics integration if the resulting grid is not sufficiently fine. Alternatively, one could use interpolation methods to estimate the state at quadrature points, but this introduces additional approximation errors.

6.3.1.4. Example: Life-Cycle Model#

The life-cycle model of consumption is a problem in economics regarding how individuals allocate resources over their lifetime, trading-off present and future consumption depending on their income and ability to save or borrow. This model demonstrates the idea that individuals tend to prefer stable consumption, even when their income varies: a behavior called “consumption smoothing”. This problem can be represented mathematically as follows:

\[\begin{split} \begin{align*} \max_{c(t)} \quad & \int_0^T e^{-\rho t} u(c(t)) dt \\ \text{subject to} \quad & \dot{A}(t) = f(A(t)) + w(t) - c(t) \\ & A(0) = A(T) = 0 \end{align*} \end{split}\]

where \([0, T]\) is the lifespan of an individual. In this model, we typically use a Constant Relative Risk Aversion (CRRA) utility function, \(u(c) = \frac{c^{1-\gamma}}{1-\gamma}\), where larger values of the parameter \(\gamma\) encode a stronger preference for a stable consumption.

The budget constraint, \(\dot{A}(t) = f(A(t)) + w(t) - c(t)\), describes the evolution of assets, \(A(t)\). Here, \(f(A(t))\) represents returns on investments, \(w(t)\) is wage income, and \(c(t)\) is consumption. The asset return function \(f(A) = 0.03A\) models a constant 3% return on investments. In our specific implementation, the choice of the wage function \(w(t) = 1 + 0.1t - 0.001t^2\) is meant to represent a career trajectory where income rises initially and then falls. The boundary conditions \(A(0) = A(T) = 0\) finally encodes the fact that individuals start and end life with zero assets (ie. no inheritances).

6.3.1.4.1. Direct Single Shooting Solution#

To solve this problem numerically, we parameterize the entire consumption path as a cubic polynomial turning our original problem into:

\[\begin{split} \begin{align*} \min_{\theta_0,\theta_1,\theta_2,\theta_3} \quad & \int_0^T e^{-\rho t} u(c(t)) dt \\ \text{subject to} \quad & \dot{A}(t) = f(A(t)) + w(t) - c(t) \\ & c(t) = \theta_0 + \theta_1t + \theta_2t^2 + \theta_3t^3 \\ & A(0) = A(T) = 0 \end{align*} \end{split}\]

We then transcribe the problem using a fourth-order Runge-Kutta method (RK4) to simulate the assets dynamics and we discretize integral cost function using:

\[\int_0^T e^{-\rho t} u(c(t)) dt \approx \sum_{i=0}^{N-1} e^{-\rho t_i} u(c(t_i)) \Delta t\]

Finally, we explicitly enforce the boundary condition \(A(T) = 0\) as an equality constraint in our optimization problem, which we solve using scipy.optimize.minimize. When using single shooting to eliminate the dynamics constraints, we the obtain the following NLP:

\[\begin{split} \begin{align*} \min_{\theta_0,\theta_1,\theta_2,\theta_3} \quad & -\sum_{i=0}^{N-1} e^{-\rho t_i} u(c_i) \Delta t \\ \text{subject to} \quad & \Phi(\theta_0,\theta_1,\theta_2,\theta_3) = 0 \\ \text{where}\quad & c_i = \theta_0 + \theta_1t_i + \theta_2t_i^2 + \theta_3t_i^3, \quad i = 0, \ldots, N \end{align*} \end{split}\]

and \(\Phi\) is defined by:

\[\begin{split} \begin{align*} \Phi(\theta_0,\theta_1,\theta_2,\theta_3) &= \Psi_N \circ \Psi_{N-1} \circ \cdots \circ \Psi_1(0) \\ \Psi_i(A) &= A + \frac{\Delta t}{6}(k_1 + 2k_2 + 2k_3 + k_4) \\ k_1 &= f(A) + w(t_i) - c_i \\ k_2 &= f(A + \frac{\Delta t}{2}k_1) + w(t_i + \frac{\Delta t}{2}) - c_{i+1/2} \\ k_3 &= f(A + \frac{\Delta t}{2}k_2) + w(t_i + \frac{\Delta t}{2}) - c_{i+1/2} \\ k_4 &= f(A + \Delta t k_3) + w(t_i + \Delta t) - c_{i+1} \\ c_i &= \theta_0 + \theta_1t_i + \theta_2t_i^2 + \theta_3t_i^3 \\ c_{i+1/2} &= \theta_0 + \theta_1(t_i + \frac{\Delta t}{2}) + \theta_2(t_i + \frac{\Delta t}{2})^2 + \theta_3(t_i + \frac{\Delta t}{2})^3 \end{align*} \end{split}\]
Hide code cell source
import numpy as np
import matplotlib.pyplot as plt
from scipy.optimize import minimize

# Model parameters
T = 50  # Time horizon
N = 1000  # Number of time steps
dt = T / N
rho = 0.05  # Discount rate
gamma = 2  # Risk aversion parameter

# Utility function
def u(c):
    return c**(1 - gamma) / (1 - gamma)

# Wage function
def w(t):
    return 1 + 0.1 * t - 0.001 * t**2

# Asset return function
def f(A):
    return 0.03 * A  # 3% return on assets

# Consumption function
def c(t, theta):
    return np.maximum(theta[0] + theta[1]*t + theta[2]*t**2 + theta[3]*t**3, 1e-6)

# Single step of RK4 method
def rk4_step(A, t, theta):
    k1 = f(A) + w(t) - c(t, theta)
    k2 = f(A + 0.5*dt*k1) + w(t + 0.5*dt) - c(t + 0.5*dt, theta)
    k3 = f(A + 0.5*dt*k2) + w(t + 0.5*dt) - c(t + 0.5*dt, theta)
    k4 = f(A + dt*k3) + w(t + dt) - c(t + dt, theta)
    return A + (dt/6) * (k1 + 2*k2 + 2*k3 + k4)

# Phi function: propagate dynamics from t=0 to t=T
def Phi(theta):
    A = 0  # Initial assets
    t = 0
    for _ in range(N):
        A = rk4_step(A, t, theta)
        t += dt
    return A

# Objective function
def objective(theta):
    t = np.linspace(0, T, N)
    consumption = c(t, theta)
    return -np.sum(np.exp(-rho * t) * u(consumption)) * dt

# Constraint function
def constraint(theta):
    return Phi(theta)

# Optimize consumption parameters
initial_guess = [1.0, 0.0, 0.0, 0.0]
cons = {'type': 'eq', 'fun': constraint}
result = minimize(objective, initial_guess, method='SLSQP', constraints=cons)
optimal_theta = result.x

# Simulate the optimal path
t = np.linspace(0, T, N)
A = np.zeros(N)
consumption = c(t, optimal_theta)

for i in range(1, N):
    A[i] = rk4_step(A[i-1], t[i-1], optimal_theta)

print(f"Optimal consumption parameters (theta): {optimal_theta}")
print(f"Final assets: {A[-1]:.4f}")

# Plotting
plt.figure(figsize=(12, 10))

plt.subplot(2, 2, 1)
plt.plot(t, A)
plt.title('Assets over time')
plt.xlabel('Time')
plt.ylabel('Assets')

plt.subplot(2, 2, 2)
plt.plot(t, consumption)
plt.title('Consumption over time')
plt.xlabel('Time')
plt.ylabel('Consumption')

plt.subplot(2, 2, 3)
plt.plot(t, w(t))
plt.title('Wage over time')
plt.xlabel('Time')
plt.ylabel('Wage')

plt.subplot(2, 2, 4)
plt.plot(t, f(A))
plt.title('Asset returns over time')
plt.xlabel('Time')
plt.ylabel('Asset returns')

plt.tight_layout()
plt.show()
Optimal consumption parameters (theta): [ 2.83489645e+00 -3.24227834e-02  3.42151299e-04 -3.00273275e-06]
Final assets: 0.0502
_images/4f025cebbf76e143069d93fb52ff635bd28f48a376427e4e5563e30629198cb3.png

6.3.2. Direct Multiple Shooting#

While direct single shooting is conceptually simple, it can suffer from numerical instabilities, especially for long time horizons or highly nonlinear systems. This phenomenon is akin to the vanishing and exploding gradient problem in deep learning (for example when unrolling an RNN over a long sequence).

Multiple shooting addresses these issues by breaking the time interval into multiple segments, treated individually indivual single shooting problems, and stiched back together via equality constraints. Not only this approach improves numerical stability but it also allows for easier incorporation of path constraints (which was otherwise difficult to do with single shooting). Another important benefit is that each subproblem can be solved in parallel, which can be beneficial from a computational point of view.

Without loss of generality, consider a Mayer problem:

\[\begin{split} \begin{aligned} \text{minimize} \quad & c(\mathbf{x}(t_f)) \\ \text{subject to} \quad & \dot{\mathbf{x}}(t) = \mathbf{f}(\mathbf{x}(t), \mathbf{u}(t)) \\ & \mathbf{g}(\mathbf{x}(t), \mathbf{u}(t)) \leq \mathbf{0} \\ & \mathbf{u}{\text{min}} \leq \mathbf{u}(t) \leq \mathbf{u}{\text{max}} \\ \text{given} \quad & \mathbf{x}(t_0) = \mathbf{x}_0 \enspace . \end{aligned} \end{split}\]

In multiple shooting, we divide the time interval \([t_0, t_f]\) into \(M\) subintervals: \([t_0, t_1], [t_1, t_2], ..., [t_{M-1}, t_M]\), where \(t_M = t_f\). For each subinterval, we introduce new optimization variables \(\mathbf{s}_i\) representing the initial state for that interval. The idea is that we want to make sure that the state at the end of an interval matches the initial state of the next one.

Heat Exchanger

The multiple shooting transcription of the problem becomes:

\[\begin{split} \begin{aligned} \underset{\boldsymbol{\theta}, \mathbf{s}_1, ..., \mathbf{s}_M}{\text{minimize}} \quad & c(\mathbf{s}_M) \\ \text{subject to} \quad & \mathbf{s}_{i+1} = \Phi_i(\boldsymbol{\theta}; \mathbf{s}_i), \quad i = 0, ..., M-1 \\ & \mathbf{g}(\mathbf{s}_i, \mathbf{u}(t_i; \boldsymbol{\theta})) \leq \mathbf{0}, \quad i = 0, ..., M \\ & \mathbf{u}_{\text{min}} \leq \mathbf{u}(t; \boldsymbol{\theta}) \leq \mathbf{u}_{\text{max}}, \quad t \in [t_0, t_f] \\ \text{given} \quad &\mathbf{s}_0 = \mathbf{x}_0 \end{aligned} \end{split}\]

Here \(\Phi_i(\boldsymbol{\theta}; \mathbf{s}_i)\) represents the final state obtained by integrating the system dynamics from \(\mathbf{s}_i\) over the interval \([t_i, t_{i+1}]\) using the parameterized control \(\mathbf{u}(t; \boldsymbol{\theta})\). The equality constraints \(\mathbf{s}_{i+1} = \Phi_i(\boldsymbol{\theta}; \mathbf{s}_i)\) ensure continuity of the state trajectory across the subintervals. These are often called “defect constraints” or “matching conditions”.

6.3.3. Direct Collocation#

While the single shooting method we discussed earlier eliminates the dynamics constraints through forward integration, an alternative approach is to keep these constraints explicit and solve the problem in a simultaneous manner. This approach, known as direct collocation, is part of a broader class of simultaneous methods in optimal control.

In direct collocation, instead of integrating the system forward in time, we discretize both the state and control trajectories and introduce additional constraints to ensure that the discretized trajectories satisfy the system dynamics. This method transforms the original continuous-time optimal control problem into a large-scale nonlinear programming (NLP) problem. Let’s consider our original Bolza problem:

\[\begin{split} \begin{aligned} \text{minimize} \quad & c(\mathbf{x}(t_f)) + \int_{t_0}^{t_f} c(\mathbf{x}(t), \mathbf{u}(t)) , dt \\ \text{subject to} \quad & \dot{\mathbf{x}}(t) = \mathbf{f}(\mathbf{x}(t), \mathbf{u}(t)) \\ & \mathbf{g}(\mathbf{x}(t), \mathbf{u}(t)) \leq \mathbf{0} \\ & \mathbf{u}_{\text{min}} \leq \mathbf{u}(t) \leq \mathbf{u}_{\text{max}} \\ \text{given} \quad & \mathbf{x}(t_0) = \mathbf{x}_0 \enspace . \end{aligned} \end{split}\]

In the direct collocation approach, we discretize the time domain into \(N\) intervals: \(t_0 < t_1 < ... < t_N = t_f\). At each node \(i\), we introduce decision variables for both the state \(\mathbf{x}_i\) and control \(\mathbf{u}_i\). The continuous dynamics are then approximated using collocation methods, which enforce the system equations at specific collocation points within each interval. The resulting NLP problem takes the form:

\[\begin{split} \begin{aligned} \underset{\mathbf{x}_0,\ldots,\mathbf{x}_N,\mathbf{u}_0,\ldots,\mathbf{u}_{N-1}}{\text{minimize}} \quad & c(\mathbf{x}_N) + \sum_{i=0}^{N-1} w_i c(\mathbf{x}_i, \mathbf{u}_i) \\ \text{subject to} \quad & \mathbf{x}_{i+1} = \mathbf{x}_i + h_i \sum_{j=1}^{k} b_j \mathbf{f}(\mathbf{x}_i^j, \mathbf{u}_i^j), \quad i = 0,...,N-1 \\ & \mathbf{g}(\mathbf{x}_i, \mathbf{u}_i) \leq \mathbf{0}, \quad i = 0,...,N-1 \\ & \mathbf{u}_{\text{min}} \leq \mathbf{u}_i \leq \mathbf{u}_{\text{max}}, \quad i = 0,...,N-1 \\ & \mathbf{x}_0 = \mathbf{x}_0^{\text{given}}\\ \end{aligned} \end{split}\]

Here, \(h_i = t_{i+1} - t_i\) is the length of the i-th interval, \(\mathbf{x}_i^j\) and \(\mathbf{u}_i^j\) are the state and control at the j-th collocation point within the i-th interval, and \(b_j\) are the weights of the collocation method. The specific form of these equations depends on the chosen collocation scheme (e.g., Hermite-Simpson, Gauss-Lobatto, etc.).

This formulation offers several advantages:

  • It provides a discretized counterpart to the original continuous-time problem.

  • It allows for path constraints on both the states and controls.

  • It often results in better numerical conditioning compared to shooting methods.

  • It allows facilitates the use of adaptive mesh refinement methods (which we don’t cover here).

In the next sections, we’ll explore specific collocation schemes and discuss their implementation in more detail.

6.3.3.1. Specific Collocation Methods#

In the following formulations:

  • \(\mathbf{x}_i\) represents the state at time \(t_i\).

  • \(\mathbf{u}_i\) represents the control at time \(t_i\).

  • \(h_i = t_{i+1} - t_i\) is the time step.

  • \(\bar{\mathbf{u}}_i\) (in RK4) represents the control at the midpoint of the interval.

  • \(\mathbf{x}_{i+\frac{1}{2}}\) and \(\mathbf{u}_{i+\frac{1}{2}}\) (in Hermite-Simpson) represent the state and control at the midpoint of the interval.

6.3.3.1.1. Euler Direct Collocation#

For the Euler method, we use linear interpolation for both state and control:

\[\begin{split} \begin{aligned} \underset{\mathbf{x}_0,\ldots,\mathbf{x}_N,\mathbf{u}_0,\ldots,\mathbf{u}_{N-1}}{\text{minimize}} \quad &c(\mathbf{x}_N) + \sum_{i=0}^{N-1} h_i c(\mathbf{x}_i, \mathbf{u}_i) \\ \text{subject to} \quad & \mathbf{x}_{i+1} = \mathbf{x}_i + h_i \mathbf{f}(\mathbf{x}_i, \mathbf{u}_i, t_i), \quad i = 0,\ldots,N-1 \\ & \mathbf{g}(\mathbf{x}_i, \mathbf{u}_i) \leq \mathbf{0}, \quad i = 0,\ldots,N-1 \\ & \mathbf{u}_{\text{min}} \leq \mathbf{u}_i \leq \mathbf{u}_{\text{max}}, \quad i = 0,\ldots,N-1 \\ & \mathbf{x}_0 = \mathbf{x}_0^{\text{given}} \end{aligned} \end{split}\]

Note that the integral of the cost function is approximated using the rectangle rule:

\[\int_{t_0}^{t_f} c(\mathbf{x}(t), \mathbf{u}(t)) dt \approx \sum_{i=0}^{N-1} h_i c(\mathbf{x}_i, \mathbf{u}_i)\]

This approximation matches the Euler integration scheme used for the dynamics:

\[\int_{t_i}^{t_{i+1}} \mathbf{f}(\mathbf{x}(t), \mathbf{u}(t), t) dt \approx h_i \mathbf{f}(\mathbf{x}_i, \mathbf{u}_i, t_i)\]

After solving the optimization problem, we obtain discrete values for the control inputs \(\mathbf{u}_i\) at each time point \(t_i\). To reconstruct the continuous control function \(\mathbf{u}(t)\), we use linear interpolation between these points. For \(t \in [t_i, t{i+1}]\), we can express \(\mathbf{u}(t)\) as:

\[\mathbf{u}(t) = \mathbf{u}_i + \frac{\mathbf{u}_{i+1} - \mathbf{u}_i}{h_i}(t - t_i)\]

This piecewise linear function provides a continuous control signal that matches the discrete optimal values found by the optimization program.

6.3.3.1.2. Trapezoidal Direct Collocation#

We use the same linear interpolation as in the Euler method, but now we enforce the dynamics at both ends of the interval:

\[\begin{split} \begin{aligned} \underset{\mathbf{x}_0,\ldots,\mathbf{x}_N,\mathbf{u}_0,\ldots,\mathbf{u}_{N}}{\text{minimize}} \quad & c(\mathbf{x}_N) + \sum_{i=0}^{N-1} \frac{h_i}{2} \left[c(\mathbf{x}_i, \mathbf{u}_i) + c(\mathbf{x}_{i+1}, \mathbf{u}_{i+1})\right] \\ \text{subject to} \quad & \mathbf{x}_{i+1} = \mathbf{x}_i + \frac{h_i}{2} \left[\mathbf{f}(\mathbf{x}_i, \mathbf{u}_i, t_i) + \mathbf{f}(\mathbf{x}_{i+1}, \mathbf{u}_{i+1}, t_{i+1})\right], \quad i = 0,\ldots,N-1 \\ & \mathbf{g}(\mathbf{x}_i, \mathbf{u}_i) \leq \mathbf{0}, \quad i = 0,\ldots,N \\ & \mathbf{u}_{\text{min}} \leq \mathbf{u}_i \leq \mathbf{u}_{\text{max}}, \quad i = 0,\ldots,N \\ & \mathbf{x}_0 = \mathbf{x}_0^{\text{given}} \end{aligned} \end{split}\]

The integral of the cost function is approximated using the trapezoidal rule:

\[\int_{t_0}^{t_f} c(\mathbf{x}(t), \mathbf{u}(t)) dt \approx \sum_{i=0}^{N-1} \frac{h_i}{2} [c(\mathbf{x}_i, \mathbf{u}_i) + c(\mathbf{x}{i+1}, \mathbf{u}{i+1})]\]

This approximation matches the trapezoidal integration scheme used for the dynamics:

\[\int_{t_i}^{t_{i+1}} \mathbf{f}(\mathbf{x}(t), \mathbf{u}(t), t) dt \approx \frac{h_i}{2} [\mathbf{f}(\mathbf{x}_i, \mathbf{u}_i, t_i) + \mathbf{f}(\mathbf{x}{i+1}, \mathbf{u}{i+1}, t_{i+1})]\]

Similar to the Euler method, after solving the optimization problem, we obtain discrete values for the control inputs \(\mathbf{u}_i\) at each time point \(t_i\). The continuous control function \(\mathbf{u}(t)\) is again reconstructed using linear interpolation. For \(t \in [t_i, t{i+1}]\):

\[\mathbf{u}(t) = \mathbf{u}_i + \frac{\mathbf{u}_{i+1} - \mathbf{u}_i}{h_i}(t - t_i)\]
6.3.3.1.3. Hermite-Simpson Direct Collocation#

For Hermite-Simpson, we use cubic interpolation for the state and quadratic for the control. The cubic interpolation comes from the fact that we use Hermite interpolation, which uses cubic polynomials to interpolate between points while matching both function values and derivatives at the endpoints.

\[\begin{split} \begin{aligned} \underset{\mathbf{x}_0,\ldots,\mathbf{x}_N, \mathbf{u}_0,\ldots,\mathbf{u}_N, \mathbf{x}_{\frac{1}{2}},\ldots,\mathbf{x}_{N-\frac{1}{2}}, \mathbf{u}_{\frac{1}{2}},\ldots,\mathbf{u}_{N-\frac{1}{2}}}{\text{minimize}} \quad &c(\mathbf{x}_N) + \sum_{i=0}^{N-1} \frac{h_i}{6} \left[c(\mathbf{x}_i, \mathbf{u}_i) + 4c(\mathbf{x}_{i+\frac{1}{2}}, \mathbf{u}_{i+\frac{1}{2}}) + c(\mathbf{x}_{i+1}, \mathbf{u}_{i+1})\right] \\ \text{subject to} \quad & \mathbf{x}_{i+1} = \mathbf{x}_i + \frac{h_i}{6} \left[\mathbf{f}(\mathbf{x}_i, \mathbf{u}_i, t_i) + 4\mathbf{f}(\mathbf{x}_{i+\frac{1}{2}}, \mathbf{u}_{i+\frac{1}{2}}, t_{i+\frac{1}{2}}) + \mathbf{f}(\mathbf{x}_{i+1}, \mathbf{u}_{i+1}, t_{i+1})\right], \quad i = 0,\ldots,N-1 \\ & \mathbf{x}_{i+\frac{1}{2}} = \frac{1}{2}(\mathbf{x}_i + \mathbf{x}_{i+1}) + \frac{h_i}{8} \left[\mathbf{f}(\mathbf{x}_i, \mathbf{u}_i, t_i) - \mathbf{f}(\mathbf{x}_{i+1}, \mathbf{u}_{i+1}, t_{i+1})\right], \quad i = 0,\ldots,N-1 \\ & \mathbf{g}(\mathbf{x}_i, \mathbf{u}_i) \leq \mathbf{0}, \quad i = 0,\ldots,N \\ & \mathbf{g}(\mathbf{x}_{i+\frac{1}{2}}, \mathbf{u}_{i+\frac{1}{2}}) \leq \mathbf{0}, \quad i = 0,\ldots,N-1 \\ & \mathbf{u}_{\text{min}} \leq \mathbf{u}_i \leq \mathbf{u}_{\text{max}}, \quad i = 0,\ldots,N \\ & \mathbf{u}_{\text{min}} \leq \mathbf{u}_{i+\frac{1}{2}} \leq \mathbf{u}_{\text{max}}, \quad i = 0,\ldots,N-1 \\ & \mathbf{x}_0 = \mathbf{x}_0^{\text{given}} \end{aligned} \end{split}\]

The objective function uses Simpson’s rule to approximate the integral of the cost:

\[\sum_{i=0}^{N-1} \frac{h_i}{6} \left[c(\mathbf{x}_i, \mathbf{u}_i) + 4c(\mathbf{x}_{i+\frac{1}{2}}, \mathbf{u}_{i+\frac{1}{2}}) + c(\mathbf{x}_{i+1}, \mathbf{u}_{i+1})\right]\]

This is the classic Simpson’s rule formula, where the integrand is evaluated at the endpoints and midpoint of each interval. The same Simpson’s rule approximation is used in the dynamics constraint:

\[\mathbf{x}_{i+1} = \mathbf{x}_i + \frac{h_i}{6} \left[\mathbf{f}(\mathbf{x}_i, \mathbf{u}_i, t_i) + 4\mathbf{f}(\mathbf{x}_{i+\frac{1}{2}}, \mathbf{u}_{i+\frac{1}{2}}, t_{i+\frac{1}{2}}) + \mathbf{f}(\mathbf{x}_{i+1}, \mathbf{u}_{i+1}, t_{i+1})\right]\]

In other words, if you were to write the dynamics constraint in its representation through the fundamental theorem of calculus and apply Simpson’s rule, you would obtain the above representation, which happens to correspond to Hermite interpolation: hence the name Hermite-Simpson.

Once the optimization problem is solved, we obtain discrete values for the control inputs \(\mathbf{u}_i\), \(\mathbf{u}_{i+\frac{1}{2}}\), and \(\mathbf{u}_{i+1}\) for each interval \([t_i, t{i+1}]\). To reconstruct the continuous control function \(\mathbf{u}(t)\), we use quadratic interpolation within each interval. For \(t \in [t_i, t_{i+1}]\), we can express \(\mathbf{u}(t)\) as:

\[\mathbf{u}(t) = a_i + b_i(t-t_i) + c_i(t-t_i)^2\]

where the coefficients \(a_i\), \(b_i\), and \(c_i\) are determined by solving the system:

\[\begin{split} \begin{aligned} \mathbf{u}(t_i) &= \mathbf{u}_i = a_i \\ \mathbf{u}(t_i + \frac{h_i}{2}) &= \mathbf{u}_{i+\frac{1}{2}} = a_i + b_i\frac{h_i}{2} + c_i\frac{h_i^2}{4} \\ \mathbf{u}(t_{i+1}) &= \mathbf{u}_{i+1} = a_i + b_ih_i + c_ih_i^2 \end{aligned} \end{split}\]
6.3.3.1.4. Runge-Kutta 4 Direct Collocation#

For RK4, we use linear interpolation for the state and piecewise constant for the control:

\[\begin{split} \begin{aligned} \underset{\mathbf{x}_0,\ldots,\mathbf{x}_N,\mathbf{u}_0,\ldots,\mathbf{u}_{N}, \bar{\mathbf{u}}_0,\ldots,\bar{\mathbf{u}}_{N-1}}{\text{minimize}} \quad & c(\mathbf{x}_N) + \sum_{i=0}^{N-1} h_i c(\mathbf{x}_i, \mathbf{u}_i) \\ \text{subject to} \quad & \mathbf{x}_{i+1} = \mathbf{x}_i + \frac{1}{6}(\mathbf{k}_1 + 2\mathbf{k}_2 + 2\mathbf{k}_3 + \mathbf{k}_4), \quad i = 0,\ldots,N-1 \\ & \mathbf{k}_1 = h_i \mathbf{f}(\mathbf{x}_i, \mathbf{u}_i, t_i) \\ & \mathbf{k}_2 = h_i \mathbf{f}(\mathbf{x}_i + \frac{1}{2}\mathbf{k}_1, \bar{\mathbf{u}}_i, t_i + \frac{h_i}{2}) \\ & \mathbf{k}_3 = h_i \mathbf{f}(\mathbf{x}_i + \frac{1}{2}\mathbf{k}_2, \bar{\mathbf{u}}_i, t_i + \frac{h_i}{2}) \\ & \mathbf{k}_4 = h_i \mathbf{f}(\mathbf{x}_i + \mathbf{k}_3, \mathbf{u}_{i+1}, t_{i+1}) \\ & \mathbf{g}(\mathbf{x}_i, \mathbf{u}_i) \leq \mathbf{0}, \quad i = 0,\ldots,N \\ & \mathbf{u}_{\text{min}} \leq \mathbf{u}_i \leq \mathbf{u}_{\text{max}}, \quad i = 0,\ldots,N \\ & \mathbf{u}_{\text{min}} \leq \bar{\mathbf{u}}_i \leq \mathbf{u}_{\text{max}}, \quad i = 0,\ldots,N-1 \\ & \mathbf{x}_0 = \mathbf{x}_0^{\text{given}} \end{aligned} \end{split}\]

The integral of the cost function is approximated using the rectangle rule:

\[\int_{t_0}^{t_f} c(\mathbf{x}(t), \mathbf{u}(t)) dt \approx \sum_{i=0}^{N-1} h_i c(\mathbf{x}_i, \mathbf{u}_i)\]

While this is a lower-order approximation compared to the RK4 scheme used for the dynamics, it uses the same points (\(t_i\)) as the start of each RK4 step. Alternatively, we could also use a Simpson’s rule-like, which would be more consistent with the RK4 scheme and the piecewise constant control reconstruction, but also more expensive.

Given the discrete values for the control inputs \(\mathbf{u}_i\) at each time point \(t_i\), and additional midpoint controls \(\bar{\mathbf{u}}_i\), we finally reconstruct control function \(\mathbf{u}(t)\), using a piecewise constant approach with a switch at the midpoint. For \(t \in [t_i, t{i+1}]\):

\[\begin{split}\mathbf{u}(t) = \begin{cases} \mathbf{u}_i & \text{if } t_i \leq t < t_i + \frac{h_i}{2} \\ \bar{\mathbf{u}}_i & \text{if } t_i + \frac{h_i}{2} \leq t < t_{i+1} \end{cases}\end{split}\]

6.3.3.2. Example: Compressor Surge Problem#

Compressors are mechanical devices used to increase the pressure of a gas by reducing its volume. They are found in many industrial settings, from natural gas pipelines to jet engines. However, compressors can suffer from a dangerous phenomenon called “surge” when the gas flow through the compressor falls too much below its design capacity. This can happen under different circumstances such as:

  • In a natural gas pipeline system, when there is less customer demand (e.g., during warm weather when less heating is needed) the flow through the compressor lowers.

  • In a jet engine, when the pilot reduces thrust during landing, less air flows through the engine’s compressors.

  • In factory, the compressor might be connected through some equipment downstream via a valve. Closing it partially restricts gas flow, similar to pinching a garden hose, and can lead to compressor surge.

As the gas flow decreases, the compressor must work harder to maintain a steady flow. If the flow becomes too low, it can lead to a “breakdown”: a phenomenon similar to an airplane stalling at low speeds or high angles of attack. In a compressor, when this breakdown occurs the gas briefly flows backward instead of moving forward, which in turns can cause violent oscillations in pressure which can damage the compressor and the equipments depending on it. One way to address this problem is by installing a close-coupled valve (CCV), which is a device connected at the output of the compressor to quickly modulate the flow. Our aim is not to devise a optimal control approach to ensure that the compressor does not experience a surge by operating this CCV appropriately.

Following [15] and [14], we model the compressor using a simplified second-order representation:

\[\begin{split} \begin{aligned} \dot{x}_1 &= B(\Psi_e(x_1) - x_2 - u) \\ \dot{x}_2 &= \frac{1}{B}(x_1 - \Phi(x_2)) \end{aligned} \end{split}\]

Here, \(\mathbf{x} = [x_1, x_2]^T\) represents the state variables:

  • \(x_1\) is the normalized mass flow through the compressor.

  • \(x_2\) is the normalized pressure ratio across the compressor.

The control input \(u\) denotes the normalized mass flow through a CCV. The functions \(\Psi_e(x_1)\) and \(\Phi(x_2)\) represent the characteristics of the compressor and valve, respectively:

\[\begin{split} \begin{aligned} \Psi_e(x_1) &= \psi_{c0} + H\left(1 + 1.5\left(\frac{x_1}{W} - 1\right) - 0.5\left(\frac{x_1}{W} - 1\right)^3\right) \\ \Phi(x_2) &= \gamma \operatorname{sign}(x_2) \sqrt{|x_2|} \end{aligned} \end{split}\]

The system parameters are given as \(\gamma = 0.5\), \(B = 1\), \(H = 0.18\), \(\psi_{c0} = 0.3\), and \(W = 0.25\).

One possible way to pose the problem [14] is by penalizing for the deviations to the setpoints using a quadratic penalty in the instantenous cost function as well as in the terminal one. Furthermore, we also penalize for taking large actions (which are energy hungry, and potentially unsafe) within the integral term. The idea of penalzing for deviations throughout is natural way of posing the problem when solving it via single shooting. Another alternative which we will explore below is to set the desired setpoint as a hard terminal constraint.

The control objective is to stabilize the system and prevent surge, formulated as a continuous-time optimal control problem (COCP) in the Bolza form:

\[\begin{split} \begin{aligned} \text{minimize} \quad & \left[ \int_0^T \alpha(\mathbf{x}(t) - \mathbf{x}^*)^T(\mathbf{x}(t) - \mathbf{x}^*) + \kappa u(t)^2 \, dt\right] + \beta(\mathbf{x}(T) - \mathbf{x}^*)^T(\mathbf{x}(T) - \mathbf{x}^*) + R v^2 \\ \text{subject to} \quad & \dot{x}_1(t) = B(\Psi_e(x_1(t)) - x_2(t) - u(t)) \\ & \dot{x}_2(t) = \frac{1}{B}(x_1(t) - \Phi(x_2(t))) \\ & u_{\text{min}} \leq u(t) \leq u_{\text{max}} \\ & -x_2(t) + 0.4 \leq v \\ & -v \leq 0 \\ & \mathbf{x}(0) = \mathbf{x}_0 \end{aligned} \end{split}\]

The parameters \(\alpha\), \(\beta\), \(\kappa\), and \(R\) are non-negative weights that allow the designer to prioritize different aspects of performance (e.g., tight setpoint tracking vs. smooth control actions). We also constraint the control input to be within \(0 \leq u(t) \leq 0.3\) due to the physical limitations of the valve.

The authors in [14] also add a soft path constraint \(x_2(t) \geq 0.4\) to ensure that we maintain a minimum pressure at all time. This is implemented as a soft constraint using slack variables. The reason that we have the term \(R v^2\) in the objective is to penalizes violations of the soft constraint: we allow for deviations, but don’t want to do it too much.

In the experiment below, we choose the setpoint \(\mathbf{x}^* = [0.40, 0.60]^T\) as it corresponds to to an unstable equilibrium point. If we were to run the system without applying any control, we would see that the system starts to oscillate.

Hide code cell source
import numpy as np
from scipy.optimize import minimize
import matplotlib.pyplot as plt

# System parameters
gamma, B, H, psi_c0, W = 0.5, 1, 0.18, 0.3, 0.25
alpha, beta, kappa, R = 1, 0, 0.08, 0
T, N = 12, 60
dt = T / N
x1_star, x2_star = 0.40, 0.60

def psi_e(x1):
    return psi_c0 + H * (1 + 1.5 * ((x1 / W) - 1) - 0.5 * ((x1 / W) - 1)**3)

def phi(x2):
    return gamma * np.sign(x2) * np.sqrt(np.abs(x2))

def system_dynamics(x, u):
    x1, x2 = x
    dx1dt = B * (psi_e(x1) - x2 - u)
    dx2dt = (1 / B) * (x1 - phi(x2))
    return np.array([dx1dt, dx2dt])

def euler_step(x, u, dt):
    return x + dt * system_dynamics(x, u)

def instantenous_cost(x, u):
    return (alpha * np.sum((x - np.array([x1_star, x2_star]))**2) + kappa * u**2)

def terminal_cost(x):
    return beta * np.sum((x - np.array([x1_star, x2_star]))**2)

def objective_and_constraints(z):
    u, v = z[:-1], z[-1]
    x = np.zeros((N+1, 2))
    x[0] = x0
    obj = 0
    cons = []
    for i in range(N):
        x[i+1] = euler_step(x[i], u[i], dt)
        obj += dt * instantenous_cost(x[i], u[i])
        cons.append(0.4 - x[i+1, 1] - v)
    obj += terminal_cost(x[-1]) + R * v**2
    return obj, np.array(cons)

def solve_trajectory_optimization(x0, u_init):
    z0 = np.zeros(N + 1)
    z0[:-1] = u_init
    bounds = [(0, 0.3)] * N + [(0, None)]
    result = minimize(
        lambda z: objective_and_constraints(z)[0],
        z0,
        method='SLSQP',
        bounds=bounds,
        constraints={'type': 'ineq', 'fun': lambda z: -objective_and_constraints(z)[1]},
        options={'disp': True, 'maxiter': 1000, 'ftol': 1e-6}
    )
    return result.x, result

def simulate_trajectory(x0, u):
    x = np.zeros((N+1, 2))
    x[0] = x0
    for i in range(N):
        x[i+1] = euler_step(x[i], u[i], dt)
    return x

# Run optimizations and simulations
x0 = np.array([0.25, 0.25])
t = np.linspace(0, T, N+1)

# Optimized control starting from zero
z_single_shooting, _ = solve_trajectory_optimization(x0, np.zeros(N))
u_opt_shoot, v_opt_shoot = z_single_shooting[:-1], z_single_shooting[-1]
x_opt_shoot = simulate_trajectory(x0, u_opt_shoot)

# Do-nothing control (u = 0)
u_nothing = np.zeros(N)
x_nothing = simulate_trajectory(x0, u_nothing)

# Plotting
plt.figure(figsize=(15, 20))

# State variables over time
plt.subplot(3, 1, 1)
plt.plot(t, x_opt_shoot[:, 0], label='x1 (opt from 0)')
plt.plot(t, x_opt_shoot[:, 1], label='x2 (opt from 0)')
plt.plot(t, x_nothing[:, 0], ':', label='x1 (do-nothing)')
plt.plot(t, x_nothing[:, 1], ':', label='x2 (do-nothing)')
plt.axhline(y=x1_star, color='r', linestyle='--', label='x1 setpoint')
plt.axhline(y=x2_star, color='g', linestyle='--', label='x2 setpoint')
plt.xlabel('Time')
plt.ylabel('State variables')
plt.title('State variables over time')
plt.legend()
plt.grid(True)

# Phase portrait
plt.subplot(3, 1, 2)
plt.plot(x_opt_shoot[:, 0], x_opt_shoot[:, 1], label='Optimized from 0')
plt.plot(x_nothing[:, 0], x_nothing[:, 1], ':', label='Do-nothing')
plt.plot(x1_star, x2_star, 'r*', markersize=10, label='Setpoint')
plt.xlabel('x1 (mass flow)')
plt.ylabel('x2 (pressure)')
plt.title('Phase portrait')
plt.legend()
plt.grid(True)

# Control inputs
plt.subplot(3, 1, 3)
plt.plot(t[:-1], u_opt_shoot, label='Optimized from 0')
plt.plot(t[:-1], u_nothing, ':', label='Do-nothing')
plt.xlabel('Time')
plt.ylabel('Control input (u)')
plt.title('Control input over time')
plt.legend()
plt.grid(True)
Optimization terminated successfully    (Exit mode 0)
            Current function value: 0.17449650818410942
            Iterations: 37
            Function evaluations: 2306
            Gradient evaluations: 37
_images/fe88815753f0d6bbd9b63595efd3044f40e07677c85c90c7be5ff2a502d53933.png
6.3.3.2.1. Solution by Trapezoidal Collocation#

Another way to pose the problem is by imposing a terminal state constraint on the system rather than through a penalty in the integral term. In the following experiment, we use a problem formulation of the form:

\[\begin{split} \begin{aligned} \text{minimize} \quad & \left[ \int_0^T \kappa u(t)^2 \, dt\right] \\ \text{subject to} \quad & \dot{x}_1(t) = B(\Psi_e(x_1(t)) - x_2(t) - u(t)) \\ & \dot{x}_2(t) = \frac{1}{B}(x_1(t) - \Phi(x_2(t))) \\ & u_{\text{min}} \leq u(t) \leq u_{\text{max}} \\ & \mathbf{x}(0) = \mathbf{x}_0 \\ & \mathbf{x}(T) = \mathbf{x}^\star \end{aligned} \end{split}\]

We then find a control function \(u(t)\) and state trajectory \(x(t)\) using the trapezoidal collocation method.

Hide code cell source
import numpy as np
from scipy.optimize import minimize
from scipy.integrate import solve_ivp
from scipy.interpolate import interp1d
import matplotlib.pyplot as plt

# System parameters
gamma, B, H, psi_c0, W = 0.5, 1, 0.18, 0.3, 0.25
kappa = 0.08
T, N = 12, 20  # Number of collocation points
t = np.linspace(0, T, N)
dt = T / (N - 1)
x1_star, x2_star = 0.40, 0.60

def psi_e(x1):
    return psi_c0 + H * (1 + 1.5 * ((x1 / W) - 1) - 0.5 * ((x1 / W) - 1)**3)

def phi(x2):
    return gamma * np.sign(x2) * np.sqrt(np.abs(x2))

def system_dynamics(t, x, u_func):
    x1, x2 = x
    u = u_func(t)
    dx1dt = B * (psi_e(x1) - x2 - u)
    dx2dt = (1 / B) * (x1 - phi(x2))
    return [dx1dt, dx2dt]

def objective(z):
    x = z[:2*N].reshape((N, 2))
    u = z[2*N:]
    
    # Trapezoidal rule for the cost function
    cost = 0
    for i in range(N-1):
        cost += 0.5 * dt * (kappa * u[i]**2 + kappa * u[i+1]**2)
    
    return cost

def constraints(z):
    x = z[:2*N].reshape((N, 2))
    u = z[2*N:]
    
    cons = []
    
    # Dynamics constraints (trapezoidal rule)
    for i in range(N-1):
        f_i = system_dynamics(t[i], x[i], lambda t: u[i])
        f_ip1 = system_dynamics(t[i+1], x[i+1], lambda t: u[i+1])
        cons.extend(x[i+1] - x[i] - 0.5 * dt * (np.array(f_i) + np.array(f_ip1)))
    
    # Terminal constraint
    cons.extend([x[-1, 0] - x1_star, x[-1, 1] - x2_star])
    
    # Initial condition constraint
    cons.extend([x[0, 0] - x0[0], x[0, 1] - x0[1]])
    
    return np.array(cons)

def solve_trajectory_optimization(x0):
    # Initial guess
    x_init = np.linspace(x0, [x1_star, x2_star], N)
    u_init = np.zeros(N)
    z0 = np.concatenate([x_init.flatten(), u_init])
    
    # Bounds
    bounds = [(None, None)] * (2*N)  # State variables
    bounds += [(0, 0.3)] * N  # Control inputs
    
    # Constraints
    cons = {'type': 'eq', 'fun': constraints}
    
    result = minimize(
        objective,
        z0,
        method='SLSQP',
        bounds=bounds,
        constraints=cons,
        options={'disp': True, 'maxiter': 1000, 'ftol': 1e-6}
    )
    return result.x, result

# Run optimization
x0 = np.array([0.5, 0.5])
z_opt, result = solve_trajectory_optimization(x0)
x_opt_coll = z_opt[:2*N].reshape((N, 2))
u_opt_coll = z_opt[2*N:]

print(f"Optimization successful: {result.success}")
print(f"Final objective value: {result.fun}")
print(f"Final state: x1 = {x_opt_coll[-1, 0]:.4f}, x2 = {x_opt_coll[-1, 1]:.4f}")
print(f"Target state: x1 = {x1_star:.4f}, x2 = {x2_star:.4f}")

# Create interpolated control function
u_func = interp1d(t, u_opt_coll, kind='linear', bounds_error=False, fill_value=(u_opt_coll[0], u_opt_coll[-1]))

# Solve IVP with the optimized control
sol = solve_ivp(lambda t, x: system_dynamics(t, x, u_func), [0, T], x0, dense_output=True)

# Generate solution points
t_dense = np.linspace(0, T, 200)
x_ivp = sol.sol(t_dense).T

# Plotting
plt.figure(figsize=(15, 20))

# State variables over time
plt.subplot(3, 1, 1)
plt.plot(t, x_opt_coll[:, 0], 'bo-', label='x1 (collocation)')
plt.plot(t, x_opt_coll[:, 1], 'ro-', label='x2 (collocation)')
plt.plot(t_dense, x_ivp[:, 0], 'b--', label='x1 (integrated)')
plt.plot(t_dense, x_ivp[:, 1], 'r--', label='x2 (integrated)')
plt.axhline(y=x1_star, color='b', linestyle=':', label='x1 setpoint')
plt.axhline(y=x2_star, color='r', linestyle=':', label='x2 setpoint')
plt.xlabel('Time')
plt.ylabel('State variables')
plt.title('State variables over time')
plt.legend()
plt.grid(True)

# Phase portrait
plt.subplot(3, 1, 2)
plt.plot(x_opt_coll[:, 0], x_opt_coll[:, 1], 'go-', label='Collocation')
plt.plot(x_ivp[:, 0], x_ivp[:, 1], 'm--', label='Integrated')
plt.plot(x1_star, x2_star, 'r*', markersize=10, label='Setpoint')
plt.xlabel('x1 (mass flow)')
plt.ylabel('x2 (pressure)')
plt.title('Phase portrait')
plt.legend()
plt.grid(True)

# Control inputs
plt.subplot(3, 1, 3)
plt.step(t, u_opt_coll, 'g-', where='post', label='Collocation')
plt.plot(t_dense, u_func(t_dense), 'm--', label='Interpolated')
plt.xlabel('Time')
plt.ylabel('Control input (u)')
plt.title('Control input over time')
plt.legend()
plt.grid(True)

plt.tight_layout()
plt.show()
Optimization terminated successfully    (Exit mode 0)
            Current function value: 0.0023543222096489565
            Iterations: 13
            Function evaluations: 794
            Gradient evaluations: 13
Optimization successful: True
Final objective value: 0.0023543222096489565
Final state: x1 = 0.4000, x2 = 0.6000
Target state: x1 = 0.4000, x2 = 0.6000
_images/3e444c5fc262305c8ca5619db8aa7f2c913af9e20259af16018089bdefb7b584.png

You can try to vary the number of collocation points in the code an observe how the state trajectory progressively matches the ground thruth (the line denoted “integrated solution”). Note that this version of the code also lacks bound constraints on the variable \(x_2\) to ensure a minimum pressure, as we did earlier. Consider this a good exercise for you to try on your own.

6.4. System Identification#

System identification is the term used outside of machine learning communities to describe the process of inferring unknown quantities of a parameterized dynamical system from observations—in other words, learning a “world model” from data. In the simplest case, we are given a parameterized model of the system and aim to infer the values of its parameters. For example, in the compressor surge problem, one might choose to use the simplified 2nd-order model and then take measurements of the actual compressor to make the model match as closely as possible. We could measure for example the characteristics of the compressor impeller or the close-coupled valve, and then refine our models with those values.

For instance, to characterize the compressor impeller, we might vary the mass flow rate and measure the resulting pressure increase. We would then use this data to fit the compressor characteristic function \(\Psi_e(x_1)\) in our model by estimating parameters like \(\psi_{c0}\), \(H\), and \(W\). The process of planning these experiments is sometimes optimized through an Optimal Design of Experiment phase. In this problem, we aim to determine the most efficient way to collect data, typically by gathering fewer but higher-quality samples at a lower acquisition cost. While we don’t cover this material in this course, it’s worth noting its relevance to the exploration and data collection problem in machine learning. The cross-pollination of ideas between optimal design of experiments and reinforcement learning could offer valuable insights to researchers and practitioners in both fields.

6.4.1. Direct Single Shooting Approach#

A straightforward approach to system identification is to collect data from the system under varying operating conditions and find model parameters that best reconstruct the observed trajectories. This can be formulated as an L2 minimization problem, similar in structure to our optimal control problems, but with the objective function now including a summation over a discrete set of observations collected at specific (and potentially irregular) time intervals. Now, one issue we face is that the interval at which the data was collected might differ from the one used by the numerical integration procedure to simulate our model. This discrepancy would make computing the reconstruction error challenging, as the number of datapoints might differ for the two processes. There are two possible ways to address this:

By aligning the time grid used by our numerical integration procedure to match that of the data. The downside of this approach is that it may compromise the accuracy of the numerical integration, especially if the data collection intervals are irregular or too large for the system’s dynamics. By using polynomial interpolation over the numerical solution to find the missing values needed to compare with the real measurements. This approach allows us to maintain a fine, regular grid for numerical integration while still comparing against irregularly sampled data.

Mathematically, if we use this second approach with Euler integration, the problem can be expressed as:

\[\begin{split} \begin{aligned} \text{minimize}_{\boldsymbol{\theta}} \quad & \sum_{k=0}^{M-1} \|\mathbf{y}_{t_f} - \hat{\mathbf{y}}_{t_k}\|_2^2 \\ \text{subject to} \quad & \mathbf{u}_{\text{min}} \leq \mathbf{u}_i \leq \mathbf{u}_{\text{max}}, \quad i = 0, \ldots, N-1 \\ \text{where} \quad & \hat{\mathbf{y}}_{t_k} = \mathbf{h}(\mathbf{x}_{t_k}; \boldsymbol{\theta}) \\ & \mathbf{x}_{i+1} = \mathbf{x}_i + \Delta t \cdot \mathbf{f}(\mathbf{x}_i, \mathbf{u}_i; \boldsymbol{\theta}), \quad i = 0, \ldots, N-1 \\ & \mathbf{x}_{t_k} = \mathbf{x}_i + \frac{t_k - t_i}{t_{i+1} - t_i} \left(\mathbf{x}_{i+1} - \mathbf{x}_i\right), \quad t_i \leq t_k < t_{i+1} \\ \text{given} \quad & \mathbf{x}_{t_0} = \mathbf{x}_0 \\ \end{aligned} \end{split}\]

The equation for \(\mathbf{x}_{t_k}\) represents linear interpolation between the numerically integrated points to obtain \(\mathbf{x}(t_k; \boldsymbol{\theta})\) which we need to compare with the observation collected at \(t_k\).

6.4.1.1. Parameter Identification in the Compressor Surge Problem#

We consider a simple form of system identification where we will attempt to recover the value of the parameter \(B\) by reconstructing trajectories of our model and compararing with a dataset of trajectories collected on the real system.

This is just a demonstration, therefore we will pretend that the real system is the one where we set the value \(B=1\) in the 2nd order simplified model. We will then vary the initial conditions of the sytem by adding a Gaussian noise perturbation to the initial conditions with mean \(0\) and standard deviation \(0.05\). Furthermore, we will use a”do-nothing” baseline controller which we perturb with Gaussian noise with mean \(0\) and standard deviation \(0.01\).

Hide code cell source
import numpy as np
from scipy.integrate import solve_ivp
import matplotlib.pyplot as plt

# System parameters
gamma, B, H, psi_c0, W = 0.5, 1, 0.18, 0.3, 0.25

# Simulation parameters
T = 50  # Total simulation time
dt = 0.1  # Time step
t = np.arange(0, T + dt, dt)
N = len(t)

# Number of trajectories
num_trajectories = 10

def psi_e(x1):
    return psi_c0 + H * (1 + 1.5 * ((x1 / W) - 1) - 0.5 * ((x1 / W) - 1)**3)

def phi(x2):
    return gamma * np.sign(x2) * np.sqrt(np.abs(x2))

def system_dynamics(t, x, u):
    x1, x2 = x
    dx1dt = B * (psi_e(x1) - x2 - u)
    dx2dt = (1 / B) * (x1 - phi(x2))
    return [dx1dt, dx2dt]

# "Do nothing" controller with small random noise
def u_func(t):
    return np.random.normal(0, 0.01)  # Mean 0, standard deviation 0.01

# Function to simulate a single trajectory
def simulate_trajectory(x0):
    sol = solve_ivp(lambda t, x: system_dynamics(t, x, u_func(t)), [0, T], x0, t_eval=t, method='RK45')
    return sol.y[0], sol.y[1]

# Generate multiple trajectories
trajectories = []
initial_conditions = []

for i in range(num_trajectories):
    # Randomize initial conditions around [0.5, 0.5]
    x0 = np.array([0.5, 0.5]) + np.random.normal(0, 0.05, 2)
    initial_conditions.append(x0)
    x1, x2 = simulate_trajectory(x0)
    trajectories.append((x1, x2))

# Calculate control inputs (small random noise)
u = np.array([u_func(ti) for ti in t])

# Plotting
plt.figure(figsize=(15, 15))

# State variables over time
plt.subplot(3, 1, 1)
for i, (x1, x2) in enumerate(trajectories):
    plt.plot(t, x1, label=f'x1 (Traj {i+1})' if i == 0 else "_nolegend_")
    plt.plot(t, x2, label=f'x2 (Traj {i+1})' if i == 0 else "_nolegend_")
plt.xlabel('Time')
plt.ylabel('State variables')
plt.title('State variables over time (Multiple Trajectories)')
plt.legend()
plt.grid(True)

# Phase portrait
plt.subplot(3, 1, 2)
for x1, x2 in trajectories:
    plt.plot(x1, x2)
    plt.plot(x1[0], x2[0], 'bo', markersize=5)
    plt.plot(x1[-1], x2[-1], 'ro', markersize=5)
plt.xlabel('x1 (mass flow)')
plt.ylabel('x2 (pressure)')
plt.title('Phase portrait (Multiple Trajectories)')
plt.grid(True)

# Control input (small random noise)
plt.subplot(3, 1, 3)
plt.plot(t, u, 'k-')
plt.xlabel('Time')
plt.ylabel('Control input (u)')
plt.title('Control input over time (Small random noise)')
plt.grid(True)

plt.tight_layout()
plt.show()

# Save the data
np.savez('_static/compressor_surge_data_multi.npz', t=t, trajectories=trajectories, u=u, initial_conditions=initial_conditions)

print("Data collection complete. Results saved to 'compressor_surge_data_multi.npz'")
print(f"Data shape: {num_trajectories} trajectories, each with {N} time steps")
print(f"Time range: 0 to {T} seconds")
print("Initial conditions:")
for i, x0 in enumerate(initial_conditions):
    print(f"  Trajectory {i+1}: x1 = {x0[0]:.4f}, x2 = {x0[1]:.4f}")
_images/bc29e09048ebfa11d7a7aedc04162c47c63f6868eda89940f4ceaf7280f59f0e.png
Data collection complete. Results saved to 'compressor_surge_data_multi.npz'
Data shape: 10 trajectories, each with 501 time steps
Time range: 0 to 50 seconds
Initial conditions:
  Trajectory 1: x1 = 0.5024, x2 = 0.4635
  Trajectory 2: x1 = 0.5442, x2 = 0.5141
  Trajectory 3: x1 = 0.4369, x2 = 0.5356
  Trajectory 4: x1 = 0.5117, x2 = 0.4580
  Trajectory 5: x1 = 0.5337, x2 = 0.5080
  Trajectory 6: x1 = 0.5652, x2 = 0.4381
  Trajectory 7: x1 = 0.4875, x2 = 0.5705
  Trajectory 8: x1 = 0.4567, x2 = 0.5260
  Trajectory 9: x1 = 0.4676, x2 = 0.4229
  Trajectory 10: x1 = 0.4703, x2 = 0.4712

We then use this data in a direct single shooting approach where we use RK4 for numerical integration. Note that in this demonstration, the time grids align and we don’t need to do use interpolation for the state trajectories.

Hide code cell source
import numpy as np
from scipy.optimize import minimize
import matplotlib.pyplot as plt

# Load the data
data = np.load('_static/compressor_surge_data_multi.npz', allow_pickle=True)
t = data['t']
trajectories = data['trajectories']
u = data['u']
initial_conditions = data['initial_conditions']

# Known system parameters
gamma, H, psi_c0, W = 0.5, 0.18, 0.3, 0.25
# B is the parameter we want to identify
B_true = 1.0  # True value, used for comparison

def psi_e(x1):
    return psi_c0 + H * (1 + 1.5 * ((x1 / W) - 1) - 0.5 * ((x1 / W) - 1)**3)

def phi(x2):
    return gamma * np.sign(x2) * np.sqrt(np.abs(x2))

def system_dynamics(t, x, u, B):
    x1, x2 = x
    dx1dt = B * (psi_e(x1) - x2 - u)
    dx2dt = (1 / B) * (x1 - phi(x2))
    return np.array([dx1dt, dx2dt])

def rk4_step(f, t, x, u, dt, B):
    k1 = f(t, x, u, B)
    k2 = f(t + 0.5*dt, x + 0.5*dt*k1, u, B)
    k3 = f(t + 0.5*dt, x + 0.5*dt*k2, u, B)
    k4 = f(t + dt, x + dt*k3, u, B)
    return x + (dt/6) * (k1 + 2*k2 + 2*k3 + k4)

def simulate_trajectory(x0, B):
    x = np.zeros((len(t), 2))
    x[0] = x0
    for i in range(1, len(t)):
        x[i] = rk4_step(system_dynamics, t[i-1], x[i-1], u[i-1], t[i] - t[i-1], B)
    return x

def objective(B):
    error = 0
    for i, (x1_obs, x2_obs) in enumerate(trajectories):
        x_sim = simulate_trajectory(initial_conditions[i], B[0])
        error += np.sum((x_sim[:, 0] - x1_obs)**2 + (x_sim[:, 1] - x2_obs)**2)
    return error

# Perform optimization
result = minimize(objective, x0=[1.5], method='Nelder-Mead', options={'disp': True})

B_identified = result.x[0]

print(f"True B: {B_true}")
print(f"Identified B: {B_identified}")
print(f"Relative error: {abs(B_identified - B_true) / B_true * 100:.2f}%")

# Plot results
plt.figure(figsize=(15, 10))

# Plot one trajectory for comparison
traj_index = 0
x1_obs, x2_obs = trajectories[traj_index]
x_sim = simulate_trajectory(initial_conditions[traj_index], B_identified)

plt.subplot(2, 1, 1)
plt.plot(t, x1_obs, 'b-', label='Observed x1')
plt.plot(t, x2_obs, 'r-', label='Observed x2')
plt.plot(t, x_sim[:, 0], 'b--', label='Simulated x1')
plt.plot(t, x_sim[:, 1], 'r--', label='Simulated x2')
plt.xlabel('Time')
plt.ylabel('State variables')
plt.title('Observed vs Simulated Trajectory')
plt.legend()
plt.grid(True)

plt.subplot(2, 1, 2)
plt.plot(x1_obs, x2_obs, 'g-', label='Observed')
plt.plot(x_sim[:, 0], x_sim[:, 1], 'm--', label='Simulated')
plt.xlabel('x1 (mass flow)')
plt.ylabel('x2 (pressure)')
plt.title('Phase Portrait: Observed vs Simulated')
plt.legend()
plt.grid(True)

plt.tight_layout()
plt.show()
Optimization terminated successfully.
         Current function value: 9.135704
         Iterations: 15
         Function evaluations: 30
True B: 1.0
Identified B: 1.041137695312499
Relative error: 4.11%
_images/ec588b0139a5ec262d70e0d4a41a779982ce73d371b946783d10e1a8906fb5b0.png

6.4.2. Parameterization of \(f\) and Neural ODEs#

In our compressor surge problem, we were provided with a physically-motivated form for the function \(f\). This set of equations was likely derived by scientists with deep knowledge of the physical phenomena at play (i.e., gas compression). However, in complex systems, the underlying physics might not be well understood or too complicated to model explicitly. In such cases, we might opt for a more flexible, data-driven approach.

Instead of specifying a fixed structure for \(f\), we could use a “black box” model such as a neural network to learn the dynamics directly from data. The optimization problem remains conceptually the same as that of parameter identification. However, we are now optimizing over the parameters of the neural network that defines \(f\).

Another possibility is to blend the two approaches and use a grey-box model. In this approach, we typically use a physics-informed parameterization which we then supplement with a black-box model to account for the discrepancies in the observations. Mathematically, this can be expressed as:

\[ \dot{\mathbf{x}}(t) = f_{\text{physics}}(\mathbf{x}, t; \boldsymbol{\theta}_{\text{physics}}) + f_{\text{NN}}(\mathbf{x}, t; \boldsymbol{\theta}_{\text{NN}}) \]

where \(f_{\text{physics}}\) is the physics-based model with parameters \(\boldsymbol{\theta}_{\text{physics}}\), and \(f_{\text{NN}}\) is a neural network with parameters \(\boldsymbol{\theta}_{\text{NN}}\) that captures unmodeled dynamics.

We then learn the parameters of the black-box model in tandem with the output of the given physics-based model. You can think of the combination of these two models as a neural network of its own, with the key difference being that one subnetwork (the physics-based one) has frozen weights (non-adjustable parameters).

This approach is easy to implement using automatic differentiation techniques and allows us to leverage prior knowledge to make the data-driven modelling more sample efficient. From a learning perspective, it amounts to providing inductive biases to make learning more efficient and to generalize better.