Cart Pole

Cart pole is a classic problem in control theory.

from pathlib import Path
import os
import functools

from IPython.display import HTML, Image
import matplotlib.pyplot as plt
import numpy as np
from celluloid import Camera
import matplotlib.patches as patches
from scipy.integrate import odeint, solve_ivp
from simple_pid import PID

ROOT = Path("./assets/img/")

if not os.path.exists(ROOT):
    os.makedirs(ROOT)

XMIN = -12
XMAX = 12
YMIN = -10
YMAX = 10.
    
L = 3.
m = 1.
M = 3.
g = -9.81

def plot_cart_pole(x, theta, ax=None):
    
    cart_width = 3.
    cart_height = 0.7

    mass_radius = 0.1
    
    if ax is None:
        ax = plt.subplot()

    ax.set_xlim((XMIN, XMAX))
    ax.set_ylim((YMIN, YMAX))
    
    rect = patches.Rectangle((x - cart_width/2., -cart_height), cart_width, cart_height, linewidth=1, edgecolor='black', facecolor='none')
    ax.add_patch(rect)

    x_m = x + L*np.sin(theta)
    y_m = L*np.cos(theta) 
    
    ax.plot([x, x_m], [0, y_m], color="black")

    circ = patches.Circle(xy=(x_m, y_m), radius=mass_radius, edgecolor="black", facecolor="black")
    ax.add_patch(circ)
    
    fig = ax.get_figure()
    plt.close()

    png_path = ROOT / "cart_pole.png"  
    fig.savefig(png_path)
    return Image(url=png_path)

plot_cart_pole(x=0, theta=np.pi/6)

Let the pivot position be $(x, y)$, the cart mass $M$, the force on the cart$F$, the pole mass (concentrated at the end) $m$, the pole length $L$, gravity acceleration $g=-9.81$ and the angle with the vertical is $\theta$ such that the up vertical position is $\theta = 0$. The mass position is then given by:

\[\begin{align*} x_m = x + L \sin(\theta)\\\\ y_m = y + L \cos(\theta) \end{align*}\]

The mass acceleration is:

\[\begin{align*} \ddot{x}_m = \ddot{x} + L (\cos(\theta) \ddot{\theta} - \sin(\theta)\dot{\theta}^2)\\\\ \ddot{y_m} = - L (\sin(\theta)\ddot{\theta} + \cos(\theta)\dot{\theta}^2) \end{align*}\]

Using Netwon’s second law for the mass and cart:

\[\begin{align*} T \cos(\theta) + m g = m \ddot{y}_m\\\\ T \sin(\theta) = m \ddot{x}_m\\\\ F - T \sin(\theta) = M \ddot{x} \end{align*}\]

Eliminate $T$:

\[\begin{align*} F - m \ddot{x}_m = M \ddot{x}\\\\ \ddot{x}_m \cos(\theta) + g \sin(\theta) = \ddot{y}_m \sin(\theta) \end{align*}\]

Elimiate mass acceleration:

\[\begin{align*} \ddot{x} = \frac{F - mL(\cos(\theta)\ddot{\theta} - \sin(\theta) \dot{\theta}^2)}{M + m} \end{align*}\] \[\ddot{\theta} = \frac{-\ddot{x} \cos(\theta) - g \sin(\theta)}{L}\]

Eliminating $\ddot{\theta}$ from the first equation:

\[\ddot{x} = \frac{F + mg\sin(\theta)\cos(\theta) + mL\sin(\theta)\dot{\theta}^2}{M+m\sin^2(\theta)}\]

Rewrite for simulation:

\[\begin{align*} \dot{v} &= \frac{F + mg\sin(\theta)\cos(\theta) + mL\sin(\theta)\omega^2}{M+m\sin^2(\theta)} \\\\ \dot{\omega} &= \frac{-\dot{v} \cos(\theta) - g \sin(\theta)}{L}\\\\ v &= \dot{x}\\\\ \omega &= \dot{\theta} \end{align*}\]

def cart_pole_step(s, t, force_fn=None):

    if force_fn is not None:
        F = force_fn(s, t)
    else:
        F = 0.
    
    x, theta, v, omega = s

    dx = v
    dtheta = omega

    dv = (F + m*g*np.sin(theta)*np.cos(theta) + m*L*np.sin(theta)*omega**2)/(M + m*np.sin(theta)**2)

    domega = (-dv*np.cos(theta)-g*np.sin(theta))/L

    return np.array([dx, dtheta, dv, domega])

def simulate_cart_pole():

    t = np.linspace(0, 20, 200)

    # x, theta, v, omega
    s0 = np.array([-0.1, .1, 0.2, 0.5])
    
    sol = odeint(cart_pole_step, s0, t)

    fig, ax = plt.subplots(1,1)

    camera = Camera(fig)
    
    for x, theta, _, _ in sol:
        plot_cart_pole(x, theta, ax=ax)
        
        camera.snap()

    anim = camera.animate()
    plt.close()

    gif_path = ROOT / "cart_pole.gif"  
    anim.save(gif_path, writer="pillow", fps=10)
    return Image(url=gif_path)
    
simulate_cart_pole()