Cart-pole (shooting)

Cart-pole (shooting)#

In this notebook, we solve a cart-pole swing-up problem using single shooting: finding control inputs that swing a pendulum from hanging down to balancing upright.

Note

This notebook demonstrates the shooting method, where we optimize only the control trajectory and compute states by forward simulation. For contrast, see Cart-pole (collocation), which uses direct collocation to optimize both states and controls with dynamics as constraints.

Key difference from direct collocation:

  • Shooting: Optimizes only controls; states are computed by simulating dynamics forward

  • Direct collocation: Optimizes both states and controls, with dynamics enforced as equality constraints

Shooting is simpler but can struggle with long time horizons due to sensitivity of the initial value problem.

Features used:

  • Var for control trajectory (single variable containing all controls)

  • @jaxls.Cost.factory for terminal state cost and control regularization

  • Forward simulation with jax.lax.scan

Hide code cell source

 
import sys
from loguru import logger

logger.remove()
logger.add(sys.stdout, format="<level>{level: <8}</level> | {message}");

import jax
import jax.numpy as jnp
import jaxls

Cart-pole dynamics#

The cart-pole system has state \([x, \theta, \dot{x}, \dot{\theta}]\) where \(x\) is cart position and \(\theta\) is pole angle (0 = hanging down, \(\pi\) = upright). The control input is a horizontal force on the cart.

# Physical parameters.
m_cart = 1.0  # Cart mass (kg)
m_pole = 0.1  # Pole mass (kg)
length = 0.5  # Pole half-length (m)
gravity = 9.81  # Gravitational acceleration (m/s^2)

# Trajectory parameters.
n_swing = 50  # Timesteps for swing-up phase
n_hold = 10  # Timesteps to hold upright position
n_steps = n_swing + n_hold  # Total timesteps
dt = 0.05  # Time step (s)
total_time = n_steps * dt
print(f"Total trajectory time: {total_time:.2f}s ({n_swing} swing + {n_hold} hold)")

Total trajectory time: 3.00s (50 swing + 10 hold)

@jax.jit
def cart_pole_dynamics(state: jax.Array, force: jax.Array) -> jax.Array:
    """Compute state derivatives for the cart-pole system.

    State: [x, theta, x_dot, theta_dot]
    theta = 0: pole hanging down, theta = pi: pole upright

    Args:
        state: Current state [x, theta, x_dot, theta_dot]
        force: Control force [f]

    Returns:
        State derivatives [x_dot, theta_dot, x_ddot, theta_ddot]
    """
    x, theta, x_dot, theta_dot = state
    f = force[0]

    sin_th = jnp.sin(theta)
    cos_th = jnp.cos(theta)

    # Total mass.
    total_mass = m_cart + m_pole
    pole_mass_length = m_pole * length

    # Equations of motion for theta=0 being DOWN (stable equilibrium)
    temp = (f + pole_mass_length * theta_dot**2 * sin_th) / total_mass
    theta_ddot = (-gravity * sin_th - cos_th * temp) / (
        length * (4.0 / 3.0 - m_pole * cos_th**2 / total_mass)
    )
    x_ddot = temp - pole_mass_length * theta_ddot * cos_th / total_mass

    return jnp.array([x_dot, theta_dot, x_ddot, theta_ddot])


@jax.jit
def rk4_step(state: jax.Array, force: jax.Array, dt: float) -> jax.Array:
    """Runge-Kutta 4th order integration step.

    Args:
        state: Current state [x, theta, x_dot, theta_dot]
        force: Control force [f]
        dt: Time step

    Returns:
        Next state after integration
    """
    k1 = cart_pole_dynamics(state, force)
    k2 = cart_pole_dynamics(state + 0.5 * dt * k1, force)
    k3 = cart_pole_dynamics(state + 0.5 * dt * k2, force)
    k4 = cart_pole_dynamics(state + dt * k3, force)
    return state + (dt / 6.0) * (k1 + 2 * k2 + 2 * k3 + k4)

Forward simulation#

The key difference from direct collocation: we simulate the entire trajectory forward using jax.lax.scan. Given a control sequence, we compute all states by integrating the dynamics.

# Initial and target states.
initial_state = jnp.array([0.0, 0.0, 0.0, 0.0])  # x=0, theta=0 (down), zero velocity
target_state = jnp.array([0.0, jnp.pi, 0.0, 0.0])  # x=0, theta=pi (up), zero velocity


@jax.jit
def simulate_trajectory(controls: jax.Array) -> jax.Array:
    """Simulate the cart-pole system forward given a control sequence.

    Args:
        controls: Control forces, shape (n_steps, 1)

    Returns:
        State trajectory, shape (n_steps + 1, 4)
    """

    def step(state: jax.Array, control: jax.Array) -> tuple[jax.Array, jax.Array]:
        next_state = rk4_step(state, control, dt)
        return next_state, next_state

    _, trajectory = jax.lax.scan(step, initial_state, controls)
    # Prepend initial state.
    return jnp.concatenate([initial_state[None, :], trajectory], axis=0)


# Test forward simulation with zero controls.
test_controls = jnp.zeros((n_steps, 1))
test_trajectory = simulate_trajectory(test_controls)
print(f"Trajectory shape: {test_trajectory.shape}")
print(f"Final state with zero controls: {test_trajectory[-1]}")

Trajectory shape: (61, 4)
Final state with zero controls: [0. 0. 0. 0.]

Control variable and costs#

In the shooting method, we optimize a single variable containing the entire control trajectory. The costs include:

  1. Terminal cost: Penalize deviation from upright position at final time

  2. Control regularization: Penalize large control inputs

class ControlTrajectoryVar(
    jaxls.Var[jax.Array], default_factory=lambda: jnp.zeros((n_steps, 1))
):
    """Control trajectory variable: (n_steps, 1) array of forces."""


control_var = ControlTrajectoryVar(id=0)

@jaxls.Cost.factory
def terminal_cost(
    vals: jaxls.VarValues,
    var: ControlTrajectoryVar,
    target: jax.Array,
) -> jax.Array:
    """Penalize deviation from target state at final time.

    Forward simulate the trajectory, then compute error at final state.
    """
    controls = vals[var]
    trajectory = simulate_trajectory(controls)
    final_state = trajectory[-1]
    # Weight position and angle errors more heavily.
    weights = jnp.array([10.0, 20.0, 5.0, 5.0])
    return (final_state - target) * weights


@jaxls.Cost.factory
def hold_cost(
    vals: jaxls.VarValues,
    var: ControlTrajectoryVar,
    target: jax.Array,
) -> jax.Array:
    """Penalize deviation from upright during the hold phase."""
    controls = vals[var]
    trajectory = simulate_trajectory(controls)
    # Get states during hold phase (use static slice with captured n_swing)
    hold_states = jax.lax.slice_in_dim(trajectory, n_swing, n_steps + 1, axis=0)
    # Penalize deviation from target for each hold state.
    weights = jnp.array([5.0, 10.0, 2.0, 2.0])
    errors = (hold_states - target) * weights
    return errors.flatten()


@jaxls.Cost.factory
def control_regularization(
    vals: jaxls.VarValues,
    var: ControlTrajectoryVar,
) -> jax.Array:
    """Penalize control effort (L2 regularization)."""
    controls = vals[var]
    return controls.flatten() * 0.1

Solving#

costs: list[jaxls.Cost] = [
    terminal_cost(control_var, target_state),
    hold_cost(control_var, target_state),
    control_regularization(control_var),
]

print(f"Created {len(costs)} cost functions")
print(f"Control variable shape: ({n_steps}, 1)")

Created 3 cost functions
Control variable shape: (60, 1)

# Initial guess: small random controls to break symmetry.
key = jax.random.PRNGKey(42)
initial_controls = jax.random.normal(key, (n_steps, 1)) * 0.1

initial_vals = jaxls.VarValues.make([control_var.with_value(initial_controls)])

# Create the problem.
problem = jaxls.LeastSquaresProblem(costs, [control_var])

# Visualize the problem structure structure.
problem.show()

# Analyze the problem.
problem = problem.analyze()

INFO     | Building optimization problem with 3 terms and 1 variables: 3 costs, 0 eq_zero, 0 leq_zero, 0 geq_zero
INFO     | Vectorizing group with 1 costs, 1 variables each: terminal_cost
INFO     | Vectorizing group with 1 costs, 1 variables each: hold_cost
INFO     | Vectorizing group with 1 costs, 1 variables each: control_regularization

solution = problem.solve(initial_vals)

INFO     |  step #0: cost=14820.0186 lambd=0.0005 inexact_tol=1.0e-02
INFO     |      - terminal_cost(1): 3938.88013 (avg 984.72003)
INFO     |      - hold_cost(1):   10881.13477 (avg 247.29852)
INFO     |      - control_regularization(1): 0.00345 (avg 0.00006)
INFO     |      accepted=True ATb_norm=5.21e+02 cost_prev=14820.0186 cost_new=10045.2891
INFO     |  step #1: cost=10045.2891 lambd=0.0003 inexact_tol=1.0e-02
INFO     |      - terminal_cost(1): 624.77783 (avg 156.19446)
INFO     |      - hold_cost(1):   9338.08984 (avg 212.22932)
INFO     |      - control_regularization(1): 82.42189 (avg 1.37370)
INFO     |  step #2: cost=10045.2891 lambd=0.0010 inexact_tol=1.0e-02
INFO     |      - terminal_cost(1): 624.77783 (avg 156.19446)
INFO     |      - hold_cost(1):   9338.08984 (avg 212.22932)
INFO     |      - control_regularization(1): 82.42189 (avg 1.37370)
INFO     |  step #3: cost=10045.2891 lambd=0.0080 inexact_tol=1.0e-02
INFO     |      - terminal_cost(1): 624.77783 (avg 156.19446)
INFO     |      - hold_cost(1):   9338.08984 (avg 212.22932)
INFO     |      - control_regularization(1): 82.42189 (avg 1.37370)
INFO     |  step #4: cost=10045.2891 lambd=0.1280 inexact_tol=1.0e-02
INFO     |      - terminal_cost(1): 624.77783 (avg 156.19446)
INFO     |      - hold_cost(1):   9338.08984 (avg 212.22932)
INFO     |      - control_regularization(1): 82.42189 (avg 1.37370)
INFO     |      accepted=True ATb_norm=2.84e+03 cost_prev=10045.2891 cost_new=4591.7881
INFO     |  step #5: cost=4591.7881 lambd=0.0640 inexact_tol=1.0e-02
INFO     |      - terminal_cost(1): 481.48416 (avg 120.37104)
INFO     |      - hold_cost(1):   4001.11768 (avg 90.93449)
INFO     |      - control_regularization(1): 109.18585 (avg 1.81976)
INFO     |      accepted=True ATb_norm=1.30e+03 cost_prev=4591.7881 cost_new=3891.5142
INFO     |  step #6: cost=3891.5144 lambd=0.0320 inexact_tol=1.0e-02
INFO     |      - terminal_cost(1): 897.49640 (avg 224.37410)
INFO     |      - hold_cost(1):   2783.67871 (avg 63.26543)
INFO     |      - control_regularization(1): 210.33926 (avg 3.50565)
INFO     |  step #7: cost=3891.5144 lambd=0.1280 inexact_tol=1.0e-02
INFO     |      - terminal_cost(1): 897.49640 (avg 224.37410)
INFO     |      - hold_cost(1):   2783.67871 (avg 63.26543)
INFO     |      - control_regularization(1): 210.33926 (avg 3.50565)
INFO     |  step #8: cost=3891.5144 lambd=1.0240 inexact_tol=1.0e-02
INFO     |      - terminal_cost(1): 897.49640 (avg 224.37410)
INFO     |      - hold_cost(1):   2783.67871 (avg 63.26543)
INFO     |      - control_regularization(1): 210.33926 (avg 3.50565)
INFO     |      accepted=True ATb_norm=1.23e+03 cost_prev=3891.5142 cost_new=3260.1982
INFO     |  step #9: cost=3260.1980 lambd=0.5120 inexact_tol=1.0e-02
INFO     |      - terminal_cost(1): 533.23682 (avg 133.30920)
INFO     |      - hold_cost(1):   2556.04761 (avg 58.09199)
INFO     |      - control_regularization(1): 170.91367 (avg 2.84856)
INFO     |  step #10: cost=3260.1980 lambd=2.0480 inexact_tol=1.0e-02
INFO     |      - terminal_cost(1): 533.23682 (avg 133.30920)
INFO     |      - hold_cost(1):   2556.04761 (avg 58.09199)
INFO     |      - control_regularization(1): 170.91367 (avg 2.84856)
INFO     |      accepted=True ATb_norm=1.83e+03 cost_prev=3260.1982 cost_new=2329.9258
INFO     |  step #11: cost=2329.9258 lambd=1.0240 inexact_tol=1.0e-02
INFO     |      - terminal_cost(1): 276.47144 (avg 69.11786)
INFO     |      - hold_cost(1):   1897.87207 (avg 43.13346)
INFO     |      - control_regularization(1): 155.58218 (avg 2.59304)
INFO     |  step #12: cost=2329.9258 lambd=4.0960 inexact_tol=1.0e-02
INFO     |      - terminal_cost(1): 276.47144 (avg 69.11786)
INFO     |      - hold_cost(1):   1897.87207 (avg 43.13346)
INFO     |      - control_regularization(1): 155.58218 (avg 2.59304)
INFO     |      accepted=True ATb_norm=7.32e+02 cost_prev=2329.9258 cost_new=1914.3083
INFO     |  step #13: cost=1914.3085 lambd=2.0480 inexact_tol=1.0e-02
INFO     |      - terminal_cost(1): 300.81384 (avg 75.20346)
INFO     |      - hold_cost(1):   1467.84851 (avg 33.36020)
INFO     |      - control_regularization(1): 145.64615 (avg 2.42744)
INFO     |      accepted=True ATb_norm=3.53e+02 cost_prev=1914.3083 cost_new=1213.7811
INFO     |  step #14: cost=1213.7810 lambd=1.0240 inexact_tol=1.0e-02
INFO     |      - terminal_cost(1): 103.21816 (avg 25.80454)
INFO     |      - hold_cost(1):   983.81934 (avg 22.35953)
INFO     |      - control_regularization(1): 126.74356 (avg 2.11239)
INFO     |      accepted=True ATb_norm=2.27e+02 cost_prev=1213.7811 cost_new=1040.4321
INFO     |  step #15: cost=1040.4321 lambd=0.5120 inexact_tol=1.0e-02
INFO     |      - terminal_cost(1): 237.12422 (avg 59.28106)
INFO     |      - hold_cost(1):   694.68127 (avg 15.78821)
INFO     |      - control_regularization(1): 108.62659 (avg 1.81044)
INFO     |  step #16: cost=1040.4321 lambd=2.0480 inexact_tol=1.0e-02
INFO     |      - terminal_cost(1): 237.12422 (avg 59.28106)
INFO     |      - hold_cost(1):   694.68127 (avg 15.78821)
INFO     |      - control_regularization(1): 108.62659 (avg 1.81044)
INFO     |      accepted=True ATb_norm=5.61e+02 cost_prev=1040.4321 cost_new=600.8384
INFO     |  step #17: cost=600.8384 lambd=1.0240 inexact_tol=1.0e-02
INFO     |      - terminal_cost(1): 27.58143 (avg 6.89536)
INFO     |      - hold_cost(1):   470.34058 (avg 10.68956)
INFO     |      - control_regularization(1): 102.91638 (avg 1.71527)
INFO     |  step #18: cost=600.8384 lambd=4.0960 inexact_tol=1.0e-02
INFO     |      - terminal_cost(1): 27.58143 (avg 6.89536)
INFO     |      - hold_cost(1):   470.34058 (avg 10.68956)
INFO     |      - control_regularization(1): 102.91638 (avg 1.71527)
INFO     |      accepted=True ATb_norm=2.13e+02 cost_prev=600.8384 cost_new=512.1956
INFO     |  step #19: cost=512.1956 lambd=2.0480 inexact_tol=1.0e-02
INFO     |      - terminal_cost(1): 1.48645 (avg 0.37161)
INFO     |      - hold_cost(1):   410.47510 (avg 9.32898)
INFO     |      - control_regularization(1): 100.23408 (avg 1.67057)
INFO     |      accepted=True ATb_norm=4.90e+01 cost_prev=512.1956 cost_new=451.1156
INFO     |  step #20: cost=451.1155 lambd=1.0240 inexact_tol=1.0e-02
INFO     |      - terminal_cost(1): 8.34001 (avg 2.08500)
INFO     |      - hold_cost(1):   347.46051 (avg 7.89683)
INFO     |      - control_regularization(1): 95.31505 (avg 1.58858)
INFO     |      accepted=True ATb_norm=1.34e+02 cost_prev=451.1156 cost_new=397.6882
INFO     |  step #21: cost=397.6882 lambd=0.5120 inexact_tol=1.0e-02
INFO     |      - terminal_cost(1): 36.54197 (avg 9.13549)
INFO     |      - hold_cost(1):   273.99484 (avg 6.22716)
INFO     |      - control_regularization(1): 87.15141 (avg 1.45252)
INFO     |      accepted=True ATb_norm=2.91e+02 cost_prev=397.6882 cost_new=386.7931
INFO     |  step #22: cost=386.7931 lambd=0.2560 inexact_tol=1.0e-02
INFO     |      - terminal_cost(1): 99.69223 (avg 24.92306)
INFO     |      - hold_cost(1):   210.61243 (avg 4.78665)
INFO     |      - control_regularization(1): 76.48843 (avg 1.27481)
INFO     |      accepted=True ATb_norm=5.02e+02 cost_prev=386.7931 cost_new=369.3773
INFO     |  step #23: cost=369.3773 lambd=0.1280 inexact_tol=1.0e-02
INFO     |      - terminal_cost(1): 148.29185 (avg 37.07296)
INFO     |      - hold_cost(1):   155.13748 (avg 3.52585)
INFO     |      - control_regularization(1): 65.94801 (avg 1.09913)
INFO     |  step #24: cost=369.3773 lambd=0.5120 inexact_tol=1.0e-02
INFO     |      - terminal_cost(1): 148.29185 (avg 37.07296)
INFO     |      - hold_cost(1):   155.13748 (avg 3.52585)
INFO     |      - control_regularization(1): 65.94801 (avg 1.09913)
INFO     |      accepted=True ATb_norm=6.39e+02 cost_prev=369.3773 cost_new=105.6115
INFO     |  step #25: cost=105.6115 lambd=0.2560 inexact_tol=1.0e-02
INFO     |      - terminal_cost(1): 3.29711 (avg 0.82428)
INFO     |      - hold_cost(1):   39.27710 (avg 0.89266)
INFO     |      - control_regularization(1): 63.03724 (avg 1.05062)
INFO     |      accepted=True ATb_norm=9.53e+01 cost_prev=105.6115 cost_new=97.8095
INFO     |  step #26: cost=97.8095 lambd=0.1280 inexact_tol=1.0e-02
INFO     |      - terminal_cost(1): 12.33471 (avg 3.08368)
INFO     |      - hold_cost(1):   26.42151 (avg 0.60049)
INFO     |      - control_regularization(1): 59.05325 (avg 0.98422)
INFO     |  step #27: cost=97.8095 lambd=0.5120 inexact_tol=1.0e-02
INFO     |      - terminal_cost(1): 12.33471 (avg 3.08368)
INFO     |      - hold_cost(1):   26.42151 (avg 0.60049)
INFO     |      - control_regularization(1): 59.05325 (avg 0.98422)
INFO     |      accepted=True ATb_norm=1.94e+02 cost_prev=97.8095 cost_new=71.1914
INFO     |  step #28: cost=71.1914 lambd=0.2560 inexact_tol=1.0e-02
INFO     |      - terminal_cost(1): 0.43552 (avg 0.10888)
INFO     |      - hold_cost(1):   13.30290 (avg 0.30234)
INFO     |      - control_regularization(1): 57.45296 (avg 0.95755)
INFO     |      accepted=True ATb_norm=3.23e+01 cost_prev=71.1914 cost_new=67.7542
INFO     |  step #29: cost=67.7542 lambd=0.1280 inexact_tol=1.0e-02
INFO     |      - terminal_cost(1): 3.06244 (avg 0.76561)
INFO     |      - hold_cost(1):   9.84872 (avg 0.22383)
INFO     |      - control_regularization(1): 54.84309 (avg 0.91405)
INFO     |  step #30: cost=67.7542 lambd=0.5120 inexact_tol=1.0e-02
INFO     |      - terminal_cost(1): 3.06244 (avg 0.76561)
INFO     |      - hold_cost(1):   9.84872 (avg 0.22383)
INFO     |      - control_regularization(1): 54.84309 (avg 0.91405)
INFO     |      accepted=True ATb_norm=9.87e+01 cost_prev=67.7542 cost_new=60.3556
INFO     |  step #31: cost=60.3556 lambd=0.2560 inexact_tol=1.0e-02
INFO     |      - terminal_cost(1): 0.20301 (avg 0.05075)
INFO     |      - hold_cost(1):   6.53582 (avg 0.14854)
INFO     |      - control_regularization(1): 53.61679 (avg 0.89361)
INFO     |      accepted=True ATb_norm=1.90e+01 cost_prev=60.3556 cost_new=58.1938
INFO     |  step #32: cost=58.1938 lambd=0.1280 inexact_tol=1.0e-02
INFO     |      - terminal_cost(1): 1.26952 (avg 0.31738)
INFO     |      - hold_cost(1):   5.53594 (avg 0.12582)
INFO     |      - control_regularization(1): 51.38834 (avg 0.85647)
INFO     |  step #33: cost=58.1938 lambd=0.5120 inexact_tol=1.0e-02
INFO     |      - terminal_cost(1): 1.26952 (avg 0.31738)
INFO     |      - hold_cost(1):   5.53594 (avg 0.12582)
INFO     |      - control_regularization(1): 51.38834 (avg 0.85647)
INFO     |      accepted=True ATb_norm=6.33e+01 cost_prev=58.1938 cost_new=54.8431
INFO     |  step #34: cost=54.8431 lambd=0.2560 inexact_tol=1.0e-02
INFO     |      - terminal_cost(1): 0.15309 (avg 0.03827)
INFO     |      - hold_cost(1):   4.37826 (avg 0.09951)
INFO     |      - control_regularization(1): 50.31177 (avg 0.83853)
INFO     |      accepted=True ATb_norm=1.39e+01 cost_prev=54.8431 cost_new=53.2063
INFO     |  step #35: cost=53.2063 lambd=0.1280 inexact_tol=1.0e-02
INFO     |      - terminal_cost(1): 0.77675 (avg 0.19419)
INFO     |      - hold_cost(1):   4.11900 (avg 0.09361)
INFO     |      - control_regularization(1): 48.31056 (avg 0.80518)
INFO     |  step #36: cost=53.2063 lambd=0.5120 inexact_tol=1.0e-02
INFO     |      - terminal_cost(1): 0.77675 (avg 0.19419)
INFO     |      - hold_cost(1):   4.11900 (avg 0.09361)
INFO     |      - control_regularization(1): 48.31056 (avg 0.80518)
INFO     |      accepted=True ATb_norm=4.89e+01 cost_prev=53.2063 cost_new=51.0624
INFO     |  step #37: cost=51.0624 lambd=0.2560 inexact_tol=1.0e-02
INFO     |      - terminal_cost(1): 0.13675 (avg 0.03419)
INFO     |      - hold_cost(1):   3.58090 (avg 0.08138)
INFO     |      - control_regularization(1): 47.34474 (avg 0.78908)
INFO     |      accepted=True ATb_norm=1.17e+01 cost_prev=51.0624 cost_new=49.7366
INFO     |  step #38: cost=49.7366 lambd=0.1280 inexact_tol=1.0e-02
INFO     |      - terminal_cost(1): 0.60969 (avg 0.15242)
INFO     |      - hold_cost(1):   3.56744 (avg 0.08108)
INFO     |      - control_regularization(1): 45.55945 (avg 0.75932)
INFO     |  step #39: cost=49.7366 lambd=0.5120 inexact_tol=1.0e-02
INFO     |      - terminal_cost(1): 0.60969 (avg 0.15242)
INFO     |      - hold_cost(1):   3.56744 (avg 0.08108)
INFO     |      - control_regularization(1): 45.55945 (avg 0.75932)
INFO     |      accepted=True ATb_norm=4.33e+01 cost_prev=49.7366 cost_new=48.0696
INFO     |  step #40: cost=48.0696 lambd=0.2560 inexact_tol=1.0e-02
INFO     |      - terminal_cost(1): 0.12940 (avg 0.03235)
INFO     |      - hold_cost(1):   3.25045 (avg 0.07387)
INFO     |      - control_regularization(1): 44.68977 (avg 0.74483)
INFO     |      accepted=True ATb_norm=1.06e+01 cost_prev=48.0696 cost_new=46.9645
INFO     |  step #41: cost=46.9645 lambd=0.1280 inexact_tol=1.0e-02
INFO     |      - terminal_cost(1): 0.53226 (avg 0.13307)
INFO     |      - hold_cost(1):   3.33001 (avg 0.07568)
INFO     |      - control_regularization(1): 43.10223 (avg 0.71837)
INFO     |  step #42: cost=46.9645 lambd=0.5120 inexact_tol=1.0e-02
INFO     |      - terminal_cost(1): 0.53226 (avg 0.13307)
INFO     |      - hold_cost(1):   3.33001 (avg 0.07568)
INFO     |      - control_regularization(1): 43.10223 (avg 0.71837)
INFO     |      accepted=True ATb_norm=4.08e+01 cost_prev=46.9645 cost_new=45.5615
INFO     |  step #43: cost=45.5615 lambd=0.2560 inexact_tol=1.0e-02
INFO     |      - terminal_cost(1): 0.12934 (avg 0.03234)
INFO     |      - hold_cost(1):   3.10468 (avg 0.07056)
INFO     |      - control_regularization(1): 42.32747 (avg 0.70546)
INFO     |      accepted=True ATb_norm=1.07e+01 cost_prev=45.5615 cost_new=44.6283
INFO     |  step #44: cost=44.6283 lambd=0.1280 inexact_tol=1.0e-02
INFO     |      - terminal_cost(1): 0.49024 (avg 0.12256)
INFO     |      - hold_cost(1):   3.21791 (avg 0.07313)
INFO     |      - control_regularization(1): 40.92015 (avg 0.68200)
INFO     |  step #45: cost=44.6283 lambd=0.5120 inexact_tol=1.0e-02
INFO     |      - terminal_cost(1): 0.49024 (avg 0.12256)
INFO     |      - hold_cost(1):   3.21791 (avg 0.07313)
INFO     |      - control_regularization(1): 40.92015 (avg 0.68200)
INFO     |      accepted=True ATb_norm=3.98e+01 cost_prev=44.6283 cost_new=43.3925
INFO     |  step #46: cost=43.3925 lambd=0.2560 inexact_tol=1.0e-02
INFO     |      - terminal_cost(1): 0.12600 (avg 0.03150)
INFO     |      - hold_cost(1):   3.04982 (avg 0.06931)
INFO     |      - control_regularization(1): 40.21673 (avg 0.67028)
INFO     |      accepted=True ATb_norm=1.02e+01 cost_prev=43.3925 cost_new=42.6236
INFO     |  step #47: cost=42.6236 lambd=0.1280 inexact_tol=1.0e-02
INFO     |      - terminal_cost(1): 0.47628 (avg 0.11907)
INFO     |      - hold_cost(1):   3.18154 (avg 0.07231)
INFO     |      - control_regularization(1): 38.96581 (avg 0.64943)
INFO     |  step #48: cost=42.6236 lambd=0.5120 inexact_tol=1.0e-02
INFO     |      - terminal_cost(1): 0.47628 (avg 0.11907)
INFO     |      - hold_cost(1):   3.18154 (avg 0.07231)
INFO     |      - control_regularization(1): 38.96581 (avg 0.64943)
INFO     |      accepted=True ATb_norm=4.00e+01 cost_prev=42.6236 cost_new=41.4925
INFO     |  step #49: cost=41.4925 lambd=0.2560 inexact_tol=1.0e-02
INFO     |      - terminal_cost(1): 0.12552 (avg 0.03138)
INFO     |      - hold_cost(1):   3.03456 (avg 0.06897)
INFO     |      - control_regularization(1): 38.33237 (avg 0.63887)
INFO     |      accepted=True ATb_norm=1.01e+01 cost_prev=41.4925 cost_new=40.8442
INFO     |  step #50: cost=40.8442 lambd=0.1280 inexact_tol=1.0e-02
INFO     |      - terminal_cost(1): 0.46028 (avg 0.11507)
INFO     |      - hold_cost(1):   3.16730 (avg 0.07198)
INFO     |      - control_regularization(1): 37.21660 (avg 0.62028)
INFO     |  step #51: cost=40.8442 lambd=0.5120 inexact_tol=1.0e-02
INFO     |      - terminal_cost(1): 0.46028 (avg 0.11507)
INFO     |      - hold_cost(1):   3.16730 (avg 0.07198)
INFO     |      - control_regularization(1): 37.21660 (avg 0.62028)
INFO     |      accepted=True ATb_norm=4.01e+01 cost_prev=40.8442 cost_new=39.8097
INFO     |  step #52: cost=39.8097 lambd=0.2560 inexact_tol=1.0e-02
INFO     |      - terminal_cost(1): 0.12928 (avg 0.03232)
INFO     |      - hold_cost(1):   3.03455 (avg 0.06897)
INFO     |      - control_regularization(1): 36.64584 (avg 0.61076)
INFO     |      accepted=True ATb_norm=1.06e+01 cost_prev=39.8097 cost_new=39.2566
INFO     |  step #53: cost=39.2566 lambd=0.1280 inexact_tol=1.0e-02
INFO     |      - terminal_cost(1): 0.44843 (avg 0.11211)
INFO     |      - hold_cost(1):   3.16067 (avg 0.07183)
INFO     |      - control_regularization(1): 35.64752 (avg 0.59413)
INFO     |  step #54: cost=39.2566 lambd=0.5120 inexact_tol=1.0e-02
INFO     |      - terminal_cost(1): 0.44843 (avg 0.11211)
INFO     |      - hold_cost(1):   3.16067 (avg 0.07183)
INFO     |      - control_regularization(1): 35.64752 (avg 0.59413)
INFO     |      accepted=True ATb_norm=4.02e+01 cost_prev=39.2566 cost_new=38.2978
INFO     |  step #55: cost=38.2978 lambd=0.2560 inexact_tol=1.0e-02
INFO     |      - terminal_cost(1): 0.12967 (avg 0.03242)
INFO     |      - hold_cost(1):   3.03547 (avg 0.06899)
INFO     |      - control_regularization(1): 35.13261 (avg 0.58554)
INFO     |      accepted=True ATb_norm=1.04e+01 cost_prev=38.2978 cost_new=37.8261
INFO     |  step #56: cost=37.8262 lambd=0.1280 inexact_tol=1.0e-02
INFO     |      - terminal_cost(1): 0.43885 (avg 0.10971)
INFO     |      - hold_cost(1):   3.15549 (avg 0.07172)
INFO     |      - control_regularization(1): 34.23181 (avg 0.57053)
INFO     |  step #57: cost=37.8262 lambd=0.5120 inexact_tol=1.0e-02
INFO     |      - terminal_cost(1): 0.43885 (avg 0.10971)
INFO     |      - hold_cost(1):   3.15549 (avg 0.07172)
INFO     |      - control_regularization(1): 34.23181 (avg 0.57053)
INFO     |      accepted=True ATb_norm=4.04e+01 cost_prev=37.8261 cost_new=36.9287
INFO     |  step #58: cost=36.9287 lambd=0.2560 inexact_tol=1.0e-02
INFO     |      - terminal_cost(1): 0.12937 (avg 0.03234)
INFO     |      - hold_cost(1):   3.03483 (avg 0.06897)
INFO     |      - control_regularization(1): 33.76454 (avg 0.56274)
INFO     |      accepted=True ATb_norm=9.96e+00 cost_prev=36.9287 cost_new=36.5273
INFO     |  step #59: cost=36.5273 lambd=0.1280 inexact_tol=1.0e-02
INFO     |      - terminal_cost(1): 0.43243 (avg 0.10811)
INFO     |      - hold_cost(1):   3.14936 (avg 0.07158)
INFO     |      - control_regularization(1): 32.94552 (avg 0.54909)
INFO     |  step #60: cost=36.5273 lambd=0.5120 inexact_tol=1.0e-02
INFO     |      - terminal_cost(1): 0.43243 (avg 0.10811)
INFO     |      - hold_cost(1):   3.14936 (avg 0.07158)
INFO     |      - control_regularization(1): 32.94552 (avg 0.54909)
INFO     |      accepted=True ATb_norm=4.08e+01 cost_prev=36.5273 cost_new=35.6825
INFO     |  step #61: cost=35.6825 lambd=0.2560 inexact_tol=1.0e-02
INFO     |      - terminal_cost(1): 0.13140 (avg 0.03285)
INFO     |      - hold_cost(1):   3.03463 (avg 0.06897)
INFO     |      - control_regularization(1): 32.51649 (avg 0.54194)
INFO     |      accepted=True ATb_norm=1.00e+01 cost_prev=35.6825 cost_new=35.3351
INFO     |  step #62: cost=35.3351 lambd=0.1280 inexact_tol=1.0e-02
INFO     |      - terminal_cost(1): 0.42599 (avg 0.10650)
INFO     |      - hold_cost(1):   3.14402 (avg 0.07146)
INFO     |      - control_regularization(1): 31.76511 (avg 0.52942)
INFO     |  step #63: cost=35.3351 lambd=0.5120 inexact_tol=1.0e-02
INFO     |      - terminal_cost(1): 0.42599 (avg 0.10650)
INFO     |      - hold_cost(1):   3.14402 (avg 0.07146)
INFO     |      - control_regularization(1): 31.76511 (avg 0.52942)
INFO     |      accepted=True ATb_norm=4.11e+01 cost_prev=35.3351 cost_new=34.5373
INFO     |  step #64: cost=34.5373 lambd=0.2560 inexact_tol=1.0e-02
INFO     |      - terminal_cost(1): 0.13394 (avg 0.03348)
INFO     |      - hold_cost(1):   3.03435 (avg 0.06896)
INFO     |      - control_regularization(1): 31.36896 (avg 0.52282)
INFO     |      accepted=True ATb_norm=1.02e+01 cost_prev=34.5373 cost_new=34.2349
INFO     |  step #65: cost=34.2349 lambd=0.1280 inexact_tol=1.0e-02
INFO     |      - terminal_cost(1): 0.42118 (avg 0.10530)
INFO     |      - hold_cost(1):   3.13936 (avg 0.07135)
INFO     |      - control_regularization(1): 30.67434 (avg 0.51124)
INFO     |  step #66: cost=34.2349 lambd=0.5120 inexact_tol=1.0e-02
INFO     |      - terminal_cost(1): 0.42118 (avg 0.10530)
INFO     |      - hold_cost(1):   3.13936 (avg 0.07135)
INFO     |      - control_regularization(1): 30.67434 (avg 0.51124)
INFO     |      accepted=True ATb_norm=4.14e+01 cost_prev=34.2349 cost_new=33.4760
INFO     |  step #67: cost=33.4760 lambd=0.2560 inexact_tol=1.0e-02
INFO     |      - terminal_cost(1): 0.13666 (avg 0.03417)
INFO     |      - hold_cost(1):   3.03265 (avg 0.06892)
INFO     |      - control_regularization(1): 30.30665 (avg 0.50511)
INFO     |      accepted=True ATb_norm=1.05e+01 cost_prev=33.4760 cost_new=33.2118
INFO     |  step #68: cost=33.2118 lambd=0.1280 inexact_tol=1.0e-02
INFO     |      - terminal_cost(1): 0.41762 (avg 0.10441)
INFO     |      - hold_cost(1):   3.13381 (avg 0.07122)
INFO     |      - control_regularization(1): 29.66041 (avg 0.49434)
INFO     |  step #69: cost=33.2118 lambd=0.5120 inexact_tol=1.0e-02
INFO     |      - terminal_cost(1): 0.41762 (avg 0.10441)
INFO     |      - hold_cost(1):   3.13381 (avg 0.07122)
INFO     |      - control_regularization(1): 29.66041 (avg 0.49434)
INFO     |      accepted=True ATb_norm=4.19e+01 cost_prev=33.2118 cost_new=32.4855
INFO     |  step #70: cost=32.4855 lambd=0.2560 inexact_tol=1.0e-02
INFO     |      - terminal_cost(1): 0.13918 (avg 0.03480)
INFO     |      - hold_cost(1):   3.02916 (avg 0.06884)
INFO     |      - control_regularization(1): 29.31718 (avg 0.48862)
INFO     |      accepted=True ATb_norm=1.07e+01 cost_prev=32.4855 cost_new=32.2542
INFO     |  step #71: cost=32.2541 lambd=0.1280 inexact_tol=1.0e-02
INFO     |      - terminal_cost(1): 0.41507 (avg 0.10377)
INFO     |      - hold_cost(1):   3.12666 (avg 0.07106)
INFO     |      - control_regularization(1): 28.71242 (avg 0.47854)
INFO     |  step #72: cost=32.2541 lambd=0.5120 inexact_tol=1.0e-02
INFO     |      - terminal_cost(1): 0.41507 (avg 0.10377)
INFO     |      - hold_cost(1):   3.12666 (avg 0.07106)
INFO     |      - control_regularization(1): 28.71242 (avg 0.47854)
INFO     |      accepted=True ATb_norm=4.24e+01 cost_prev=32.2542 cost_new=31.5561
INFO     |  step #73: cost=31.5561 lambd=0.2560 inexact_tol=1.0e-02
INFO     |      - terminal_cost(1): 0.14141 (avg 0.03535)
INFO     |      - hold_cost(1):   3.02361 (avg 0.06872)
INFO     |      - control_regularization(1): 28.39105 (avg 0.47318)
INFO     |      accepted=True ATb_norm=1.09e+01 cost_prev=31.5561 cost_new=31.3545
INFO     |  step #74: cost=31.3545 lambd=0.1280 inexact_tol=1.0e-02
INFO     |      - terminal_cost(1): 0.41428 (avg 0.10357)
INFO     |      - hold_cost(1):   3.11785 (avg 0.07086)
INFO     |      - control_regularization(1): 27.82241 (avg 0.46371)
INFO     |  step #75: cost=31.3545 lambd=0.5120 inexact_tol=1.0e-02
INFO     |      - terminal_cost(1): 0.41428 (avg 0.10357)
INFO     |      - hold_cost(1):   3.11785 (avg 0.07086)
INFO     |      - control_regularization(1): 27.82241 (avg 0.46371)
INFO     |      accepted=True ATb_norm=4.31e+01 cost_prev=31.3545 cost_new=30.6785
INFO     |  step #76: cost=30.6785 lambd=0.2560 inexact_tol=1.0e-02
INFO     |      - terminal_cost(1): 0.14342 (avg 0.03586)
INFO     |      - hold_cost(1):   3.01414 (avg 0.06850)
INFO     |      - control_regularization(1): 27.52098 (avg 0.45868)
INFO     |      accepted=True ATb_norm=1.12e+01 cost_prev=30.6785 cost_new=30.5005
INFO     |  step #77: cost=30.5005 lambd=0.1280 inexact_tol=1.0e-02
INFO     |      - terminal_cost(1): 0.41240 (avg 0.10310)
INFO     |      - hold_cost(1):   3.10431 (avg 0.07055)
INFO     |      - control_regularization(1): 26.98381 (avg 0.44973)
INFO     |  step #78: cost=30.5005 lambd=0.5120 inexact_tol=1.0e-02
INFO     |      - terminal_cost(1): 0.41240 (avg 0.10310)
INFO     |      - hold_cost(1):   3.10431 (avg 0.07055)
INFO     |      - control_regularization(1): 26.98381 (avg 0.44973)
INFO     |      accepted=True ATb_norm=4.37e+01 cost_prev=30.5005 cost_new=29.8452
INFO     |  step #79: cost=29.8452 lambd=0.2560 inexact_tol=1.0e-02
INFO     |      - terminal_cost(1): 0.14567 (avg 0.03642)
INFO     |      - hold_cost(1):   3.00107 (avg 0.06821)
INFO     |      - control_regularization(1): 26.69844 (avg 0.44497)
INFO     |      accepted=True ATb_norm=1.15e+01 cost_prev=29.8452 cost_new=29.6934
INFO     |  step #80: cost=29.6934 lambd=0.1280 inexact_tol=1.0e-02
INFO     |      - terminal_cost(1): 0.41490 (avg 0.10372)
INFO     |      - hold_cost(1):   3.08940 (avg 0.07021)
INFO     |      - control_regularization(1): 26.18906 (avg 0.43648)
INFO     |  step #81: cost=29.6934 lambd=0.5120 inexact_tol=1.0e-02
INFO     |      - terminal_cost(1): 0.41490 (avg 0.10372)
INFO     |      - hold_cost(1):   3.08940 (avg 0.07021)
INFO     |      - control_regularization(1): 26.18906 (avg 0.43648)
INFO     |      accepted=True ATb_norm=4.47e+01 cost_prev=29.6934 cost_new=29.0511
INFO     |  step #82: cost=29.0511 lambd=0.2560 inexact_tol=1.0e-02
INFO     |      - terminal_cost(1): 0.14744 (avg 0.03686)
INFO     |      - hold_cost(1):   2.98357 (avg 0.06781)
INFO     |      - control_regularization(1): 25.92013 (avg 0.43200)
INFO     |      accepted=True ATb_norm=1.18e+01 cost_prev=29.0511 cost_new=28.9216
INFO     |  step #83: cost=28.9216 lambd=0.1280 inexact_tol=1.0e-02
INFO     |      - terminal_cost(1): 0.41661 (avg 0.10415)
INFO     |      - hold_cost(1):   3.07080 (avg 0.06979)
INFO     |      - control_regularization(1): 25.43424 (avg 0.42390)
INFO     |  step #84: cost=28.9216 lambd=0.5120 inexact_tol=1.0e-02
INFO     |      - terminal_cost(1): 0.41661 (avg 0.10415)
INFO     |      - hold_cost(1):   3.07080 (avg 0.06979)
INFO     |      - control_regularization(1): 25.43424 (avg 0.42390)
INFO     |      accepted=True ATb_norm=4.56e+01 cost_prev=28.9216 cost_new=28.2908
INFO     |  step #85: cost=28.2908 lambd=0.2560 inexact_tol=1.0e-02
INFO     |      - terminal_cost(1): 0.14907 (avg 0.03727)
INFO     |      - hold_cost(1):   2.96335 (avg 0.06735)
INFO     |      - control_regularization(1): 25.17836 (avg 0.41964)
INFO     |      accepted=True ATb_norm=1.21e+01 cost_prev=28.2908 cost_new=28.1818
INFO     |  step #86: cost=28.1818 lambd=0.1280 inexact_tol=1.0e-02
INFO     |      - terminal_cost(1): 0.41847 (avg 0.10462)
INFO     |      - hold_cost(1):   3.04880 (avg 0.06929)
INFO     |      - control_regularization(1): 24.71455 (avg 0.41191)
INFO     |  step #87: cost=28.1818 lambd=0.5120 inexact_tol=1.0e-02
INFO     |      - terminal_cost(1): 0.41847 (avg 0.10462)
INFO     |      - hold_cost(1):   3.04880 (avg 0.06929)
INFO     |      - control_regularization(1): 24.71455 (avg 0.41191)
INFO     |      accepted=True ATb_norm=4.67e+01 cost_prev=28.1818 cost_new=27.5605
INFO     |  step #88: cost=27.5605 lambd=0.2560 inexact_tol=1.0e-02
INFO     |      - terminal_cost(1): 0.15038 (avg 0.03759)
INFO     |      - hold_cost(1):   2.93858 (avg 0.06679)
INFO     |      - control_regularization(1): 24.47159 (avg 0.40786)
INFO     |      accepted=True ATb_norm=1.25e+01 cost_prev=27.5605 cost_new=27.4721
INFO     |  step #89: cost=27.4721 lambd=0.1280 inexact_tol=1.0e-02
INFO     |      - terminal_cost(1): 0.42138 (avg 0.10535)
INFO     |      - hold_cost(1):   3.02270 (avg 0.06870)
INFO     |      - control_regularization(1): 24.02804 (avg 0.40047)
INFO     |  step #90: cost=27.4721 lambd=0.5120 inexact_tol=1.0e-02
INFO     |      - terminal_cost(1): 0.42138 (avg 0.10535)
INFO     |      - hold_cost(1):   3.02270 (avg 0.06870)
INFO     |      - control_regularization(1): 24.02804 (avg 0.40047)
INFO     |      accepted=True ATb_norm=4.78e+01 cost_prev=27.4721 cost_new=26.8581
INFO     |  step #91: cost=26.8581 lambd=0.2560 inexact_tol=1.0e-02
INFO     |      - terminal_cost(1): 0.15145 (avg 0.03786)
INFO     |      - hold_cost(1):   2.90931 (avg 0.06612)
INFO     |      - control_regularization(1): 23.79738 (avg 0.39662)
INFO     |      accepted=True ATb_norm=1.29e+01 cost_prev=26.8581 cost_new=26.7873
INFO     |  step #92: cost=26.7873 lambd=0.1280 inexact_tol=1.0e-02
INFO     |      - terminal_cost(1): 0.42303 (avg 0.10576)
INFO     |      - hold_cost(1):   2.99232 (avg 0.06801)
INFO     |      - control_regularization(1): 23.37196 (avg 0.38953)
INFO     |  step #93: cost=26.7873 lambd=0.5120 inexact_tol=1.0e-02
INFO     |      - terminal_cost(1): 0.42303 (avg 0.10576)
INFO     |      - hold_cost(1):   2.99232 (avg 0.06801)
INFO     |      - control_regularization(1): 23.37196 (avg 0.38953)
INFO     |      accepted=True ATb_norm=4.90e+01 cost_prev=26.7873 cost_new=26.1809
INFO     |  step #94: cost=26.1809 lambd=0.2560 inexact_tol=1.0e-02
INFO     |      - terminal_cost(1): 0.15189 (avg 0.03797)
INFO     |      - hold_cost(1):   2.87673 (avg 0.06538)
INFO     |      - control_regularization(1): 23.15228 (avg 0.38587)
INFO     |      accepted=True ATb_norm=1.32e+01 cost_prev=26.1809 cost_new=26.1282
INFO     |  step #95: cost=26.1282 lambd=0.1280 inexact_tol=1.0e-02
INFO     |      - terminal_cost(1): 0.42522 (avg 0.10630)
INFO     |      - hold_cost(1):   2.95806 (avg 0.06723)
INFO     |      - control_regularization(1): 22.74489 (avg 0.37908)
INFO     |  step #96: cost=26.1282 lambd=0.5120 inexact_tol=1.0e-02
INFO     |      - terminal_cost(1): 0.42522 (avg 0.10630)
INFO     |      - hold_cost(1):   2.95806 (avg 0.06723)
INFO     |      - control_regularization(1): 22.74489 (avg 0.37908)
INFO     |      accepted=True ATb_norm=5.02e+01 cost_prev=26.1282 cost_new=25.5271
INFO     |  step #97: cost=25.5271 lambd=0.2560 inexact_tol=1.0e-02
INFO     |      - terminal_cost(1): 0.15250 (avg 0.03813)
INFO     |      - hold_cost(1):   2.83783 (avg 0.06450)
INFO     |      - control_regularization(1): 22.53681 (avg 0.37561)
INFO     |      accepted=True ATb_norm=1.36e+01 cost_prev=25.5271 cost_new=25.4883
INFO     |  step #98: cost=25.4883 lambd=0.1280 inexact_tol=1.0e-02
INFO     |      - terminal_cost(1): 0.42504 (avg 0.10626)
INFO     |      - hold_cost(1):   2.91902 (avg 0.06634)
INFO     |      - control_regularization(1): 22.14427 (avg 0.36907)
INFO     |  step #99: cost=25.4883 lambd=0.5120 inexact_tol=1.0e-02
INFO     |      - terminal_cost(1): 0.42504 (avg 0.10626)
INFO     |      - hold_cost(1):   2.91902 (avg 0.06634)
INFO     |      - control_regularization(1): 22.14427 (avg 0.36907)
INFO     |      accepted=True ATb_norm=5.13e+01 cost_prev=25.4883 cost_new=24.8963
INFO     | Terminated @ iteration #100: cost=24.8963 criteria=[0 0 0], term_deltas=2.3e-02,1.5e+01,1.4e-02 (solved in 0.7140 sec)

Visualization#

# Extract solution.
optimal_controls = solution[control_var]
optimal_trajectory = simulate_trajectory(optimal_controls)

times = jnp.linspace(0, total_time, n_steps + 1)
control_times = jnp.linspace(0, total_time - dt, n_steps)

print(
    f"Final state: x={float(optimal_trajectory[-1, 0]):.3f}, theta={float(optimal_trajectory[-1, 1]):.3f}"
)
print(f"Target theta (pi): {jnp.pi:.3f}")
print(f"Terminal angle error: {abs(float(optimal_trajectory[-1, 1]) - jnp.pi):.4f} rad")

Final state: x=-0.003, theta=3.141
Target theta (pi): 3.142
Terminal angle error: 0.0004 rad

Hide code cell source

 
import plotly.graph_objects as go
from plotly.subplots import make_subplots
from IPython.display import HTML

# Animation of cart-pole swing-up.
cart_y_offset = 1.2  # Raise cart so pole can swing freely above ground.
wheel_radius = 0.05


def create_cart_pole_frame(x: float, theta: float) -> dict:
    """Create a single frame of cart-pole visualization.

    Args:
        x: Cart position.
        theta: Pole angle (0 = down, pi = up).

    Returns:
        Dictionary with cart, pole, mass, and wheel coordinates.
    """
    cart_width, cart_height = 0.4, 0.15
    pole_len = 2 * length  # Full pole length.
    pole_x = x + pole_len * jnp.sin(theta)
    pole_y = cart_y_offset - pole_len * jnp.cos(theta)

    # Wheel positions.
    wheel_y = cart_y_offset - cart_height / 2 - wheel_radius

    return dict(
        cart_x=[
            x - cart_width / 2,
            x + cart_width / 2,
            x + cart_width / 2,
            x - cart_width / 2,
            x - cart_width / 2,
        ],
        cart_y=[
            cart_y_offset - cart_height / 2,
            cart_y_offset - cart_height / 2,
            cart_y_offset + cart_height / 2,
            cart_y_offset + cart_height / 2,
            cart_y_offset - cart_height / 2,
        ],
        pole_x=[x, float(pole_x)],
        pole_y=[cart_y_offset, float(pole_y)],
        mass_x=[float(pole_x)],
        mass_y=[float(pole_y)],
        wheel1_x=[x - cart_width / 3],
        wheel1_y=[wheel_y],
        wheel2_x=[x + cart_width / 3],
        wheel2_y=[wheel_y],
    )


# Create animation frames (use every 2nd timestep for performance).
frame_step = 2
frame_indices = list(range(0, n_steps + 1, frame_step))
frames = []
for i in frame_indices:
    frame_data = create_cart_pole_frame(
        float(optimal_trajectory[i, 0]), float(optimal_trajectory[i, 1])
    )
    frames.append(
        go.Frame(
            data=[
                go.Scatter(
                    x=frame_data["cart_x"],
                    y=frame_data["cart_y"],
                    fill="toself",
                    fillcolor="steelblue",
                    line=dict(color="darkblue", width=2),
                ),
                go.Scatter(
                    x=frame_data["pole_x"],
                    y=frame_data["pole_y"],
                    mode="lines",
                    line=dict(color="sienna", width=8),
                ),
                go.Scatter(
                    x=frame_data["mass_x"],
                    y=frame_data["mass_y"],
                    mode="markers",
                    marker=dict(
                        size=18, color="sienna", line=dict(color="black", width=1)
                    ),
                ),
                go.Scatter(
                    x=frame_data["wheel1_x"],
                    y=frame_data["wheel1_y"],
                    mode="markers",
                    marker=dict(size=12, color="black", symbol="circle"),
                ),
                go.Scatter(
                    x=frame_data["wheel2_x"],
                    y=frame_data["wheel2_y"],
                    mode="markers",
                    marker=dict(size=12, color="black", symbol="circle"),
                ),
            ],
            name=str(i),
        )
    )

# Initial frame.
init_frame = create_cart_pole_frame(
    float(optimal_trajectory[0, 0]), float(optimal_trajectory[0, 1])
)
rail_y = cart_y_offset - 0.15 / 2 - wheel_radius * 2

fig_anim = go.Figure(
    data=[
        go.Scatter(
            x=init_frame["cart_x"],
            y=init_frame["cart_y"],
            fill="toself",
            fillcolor="steelblue",
            line=dict(color="darkblue", width=2),
            name="Cart",
        ),
        go.Scatter(
            x=init_frame["pole_x"],
            y=init_frame["pole_y"],
            mode="lines",
            line=dict(color="sienna", width=8),
            name="Pole",
        ),
        go.Scatter(
            x=init_frame["mass_x"],
            y=init_frame["mass_y"],
            mode="markers",
            marker=dict(size=18, color="sienna", line=dict(color="black", width=1)),
            showlegend=False,
        ),
        go.Scatter(
            x=init_frame["wheel1_x"],
            y=init_frame["wheel1_y"],
            mode="markers",
            marker=dict(size=12, color="black", symbol="circle"),
            showlegend=False,
        ),
        go.Scatter(
            x=init_frame["wheel2_x"],
            y=init_frame["wheel2_y"],
            mode="markers",
            marker=dict(size=12, color="black", symbol="circle"),
            showlegend=False,
        ),
        go.Scatter(
            x=[-2.5, 2.5],
            y=[rail_y, rail_y],
            mode="lines",
            line=dict(color="gray", width=6),
            name="Rail",
        ),
    ],
    frames=frames,
    layout=go.Layout(
        title="Cart-Pole (Shooting)",
        xaxis=dict(
            range=[-2.5, 2.5], title="x (m)", constrain="domain", showgrid=False
        ),
        yaxis=dict(
            range=[-0.1, 2.5],
            title="y (m)",
            scaleanchor="x",
            scaleratio=1,
            showgrid=False,
        ),
        updatemenus=[
            dict(
                type="buttons",
                showactive=False,
                y=1.15,
                x=0.5,
                xanchor="center",
                buttons=[
                    dict(
                        label="Play",
                        method="animate",
                        args=[
                            None,
                            dict(
                                frame=dict(duration=50, redraw=True),
                                fromcurrent=True,
                                transition=dict(duration=0),
                            ),
                        ],
                    ),
                    dict(
                        label="Pause",
                        method="animate",
                        args=[
                            [None],
                            dict(
                                frame=dict(duration=0, redraw=False),
                                mode="immediate",
                            ),
                        ],
                    ),
                ],
            )
        ],
        sliders=[
            dict(
                active=0,
                yanchor="top",
                xanchor="left",
                currentvalue=dict(
                    prefix="Time: ",
                    suffix="s",
                    visible=True,
                    xanchor="center",
                    offset=20,
                ),
                pad=dict(b=10, t=60),
                steps=[
                    dict(
                        args=[
                            [str(i)],
                            dict(
                                frame=dict(duration=0, redraw=True),
                                mode="immediate",
                                transition=dict(duration=0),
                            ),
                        ],
                        label=f"{float(times[i]):.2f}",
                        method="animate",
                    )
                    for i in frame_indices
                ],
                x=0.1,
                y=0,
                len=0.8,
            )
        ],
        height=500,
        showlegend=False,
        margin=dict(t=80, b=100),
        plot_bgcolor="white",
    ),
)
HTML(fig_anim.to_html(full_html=False, include_plotlyjs="cdn", auto_play=False))

Hide code cell source

 
# State and control trajectories.
fig_traj = make_subplots(
    rows=2,
    cols=2,
    subplot_titles=("Cart Position", "Pole Angle", "Velocities", "Control Force"),
    vertical_spacing=0.15,
    horizontal_spacing=0.1,
)

# Cart position.
fig_traj.add_trace(
    go.Scatter(
        x=[float(t) for t in times],
        y=[float(s) for s in optimal_trajectory[:, 0]],
        mode="lines",
        line=dict(color="steelblue", width=2),
        name="x",
    ),
    row=1,
    col=1,
)

# Pole angle.
fig_traj.add_trace(
    go.Scatter(
        x=[float(t) for t in times],
        y=[float(s) for s in optimal_trajectory[:, 1]],
        mode="lines",
        line=dict(color="coral", width=2),
        name="theta",
    ),
    row=1,
    col=2,
)
fig_traj.add_hline(
    y=float(jnp.pi),
    line_dash="dash",
    line_color="gray",
    row=1,
    col=2,
    annotation_text="pi",
)

# Velocities.
fig_traj.add_trace(
    go.Scatter(
        x=[float(t) for t in times],
        y=[float(s) for s in optimal_trajectory[:, 2]],
        mode="lines",
        line=dict(color="steelblue", width=2),
        name="x_dot",
    ),
    row=2,
    col=1,
)
fig_traj.add_trace(
    go.Scatter(
        x=[float(t) for t in times],
        y=[float(s) for s in optimal_trajectory[:, 3]],
        mode="lines",
        line=dict(color="coral", width=2, dash="dash"),
        name="theta_dot",
    ),
    row=2,
    col=1,
)

# Control force.
fig_traj.add_trace(
    go.Scatter(
        x=[float(t) for t in control_times],
        y=[float(u) for u in optimal_controls[:, 0]],
        mode="lines",
        line=dict(color="forestgreen", width=2),
        name="F",
        fill="tozeroy",
        fillcolor="rgba(34, 139, 34, 0.2)",
    ),
    row=2,
    col=2,
)

# Axis labels.
fig_traj.update_xaxes(title_text="Time (s)", row=1, col=1)
fig_traj.update_yaxes(title_text="Position (m)", row=1, col=1)
fig_traj.update_xaxes(title_text="Time (s)", row=1, col=2)
fig_traj.update_yaxes(title_text="Angle (rad)", row=1, col=2)
fig_traj.update_xaxes(title_text="Time (s)", row=2, col=1)
fig_traj.update_yaxes(title_text="Velocity", row=2, col=1)
fig_traj.update_xaxes(title_text="Time (s)", row=2, col=2)
fig_traj.update_yaxes(title_text="Force (N)", row=2, col=2)

fig_traj.update_layout(height=500, showlegend=False, margin=dict(t=40, b=40))
HTML(fig_traj.to_html(full_html=False, include_plotlyjs="cdn"))

Comparison with direct collocation#

The shooting method optimizes only the control trajectory (60 decision variables for 60 timesteps), while direct collocation in Cart-pole (collocation) optimizes both states and controls (305 decision variables: 61 states x 4 + 60 controls).

Shooting method advantages:

  • Fewer decision variables

  • Dynamics are satisfied by construction (no defect constraints)

  • Simpler problem formulation

Shooting method disadvantages:

  • Sensitivity to initial conditions grows exponentially with time horizon

  • Harder to add state constraints (must be handled through cost)

  • Can struggle with long horizons or unstable dynamics

For this cart-pole problem, both methods work well because the time horizon is short. For longer horizons or more complex systems, direct collocation often provides better convergence.

For solver configuration options, see jaxls.TrustRegionConfig and jaxls.TerminationConfig.

Contents