Cart-pole (collocation)

Cart-pole (collocation)#

In this notebook, we solve a cart-pole swing-up problem using direct collocation: an alternative to single shooting that optimizes the full state trajectory simultaneously.

Features used:

  • Var subclassing with batched IDs for state and control trajectories

  • @jaxls.Cost.factory with batched arguments

  • Equality constraints (constraint_eq_zero): dynamics collocation

  • Inequality constraints (constraint_leq_zero): control bounds

  • Augmented Lagrangian solver for constrained optimization

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)
force_limit = 10.0  # Maximum control force (N)

# 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

    Note: We use the standard cart-pole equations but with theta measured
    from hanging down, so the gravitational term has a negative sign.
    """
    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)
    # The gravitational term is negated compared to theta=0 being UP.
    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."""
    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)

Variables#

We define state variables at each timestep and control inputs between timesteps. Using batched IDs for efficient construction:

class StateVar(jaxls.Var[jax.Array], default_factory=lambda: jnp.zeros(4)):
    """State variable: [x, theta, x_dot, theta_dot]."""


class ControlVar(jaxls.Var[jax.Array], default_factory=lambda: jnp.zeros(1)):
    """Control variable: [force]."""


# Create batched variables.
state_vars = StateVar(id=jnp.arange(n_steps + 1))  # States at t=0, 1, ..., n_steps
control_vars = ControlVar(id=jnp.arange(n_steps))  # Controls at t=0, 1, ..., n_steps-1

print(f"State variables: {n_steps + 1} (shape: {state_vars.id.shape})")
print(f"Control variables: {n_steps} (shape: {control_vars.id.shape})")

State variables: 61 (shape: (61,))
Control variables: 60 (shape: (60,))

Cost functions#

We define costs for:

  1. Dynamics constraints: RK4 integration between consecutive states

  2. Boundary costs: Penalize deviation from initial and target states

  3. Control effort: Minimize force usage

  4. Control bounds: Keep force within limits

@jaxls.Cost.factory(kind="constraint_eq_zero")
def dynamics_constraint(
    vals: jaxls.VarValues,
    state_k: StateVar,
    state_k1: StateVar,
    control_k: ControlVar,
    dt: float,
) -> jax.Array:
    """Enforce RK4 dynamics between consecutive states."""
    s_k = vals[state_k]
    s_k1 = vals[state_k1]
    u_k = vals[control_k]
    s_next = rk4_step(s_k, u_k, dt)
    return s_k1 - s_next


@jaxls.Cost.factory
def boundary_cost(
    vals: jaxls.VarValues,
    var: StateVar,
    target: jax.Array,
) -> jax.Array:
    """Penalize deviation from target state."""
    return (vals[var] - target) * 100.0


@jaxls.Cost.factory
def control_cost(
    vals: jaxls.VarValues,
    var: ControlVar,
) -> jax.Array:
    """Minimize control effort."""
    return vals[var] * 1.0


@jaxls.Cost.factory(kind="constraint_leq_zero")
def control_upper_bound(
    vals: jaxls.VarValues,
    var: ControlVar,
    limit: float,
) -> jax.Array:
    """Control <= limit."""
    return vals[var] - limit


@jaxls.Cost.factory(kind="constraint_leq_zero")
def control_lower_bound(
    vals: jaxls.VarValues,
    var: ControlVar,
    limit: float,
) -> jax.Array:
    """-limit <= control, i.e., -limit - control <= 0."""
    return -limit - vals[var]

Problem construction#

Initial state: pole hanging down (\(\theta = 0\)). Target state: pole upright (\(\theta = \pi\)), centered at origin with zero velocity.

# Boundary conditions.
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

# State indices for dynamics constraints.
state_k_ids = jnp.arange(n_steps)  # 0, 1, ..., n_steps-1
state_k1_ids = jnp.arange(1, n_steps + 1)  # 1, 2, ..., n_steps
control_k_ids = jnp.arange(n_steps)  # 0, 1, ..., n_steps-1

# Hold phase: penalize states during hold to stay at target.
hold_state_ids = jnp.arange(n_swing, n_steps + 1)  # States during hold phase
n_hold_states = len(hold_state_ids)
hold_targets = jnp.tile(target_state, (n_hold_states, 1))  # Broadcast for batching

# Build costs using batched construction.
costs: list[jaxls.Cost] = [
    # Dynamics constraints (batched)
    dynamics_constraint(
        StateVar(id=state_k_ids),
        StateVar(id=state_k1_ids),
        ControlVar(id=control_k_ids),
        dt,
    ),
    # Boundary costs.
    boundary_cost(StateVar(id=0), initial_state),
    # Hold phase: penalize all hold states deviating from target.
    boundary_cost(StateVar(id=hold_state_ids), hold_targets),
    # Control effort (batched)
    control_cost(control_vars),
    # Control bounds (batched)
    control_upper_bound(control_vars, force_limit),
    control_lower_bound(control_vars, force_limit),
]

print(f"Created {len(costs)} batched cost objects")

Created 6 batched cost objects

# Initial guess: linear interpolation for states, zero controls.
t_interp = jnp.linspace(0, 1, n_steps + 1)[:, None]
initial_states = initial_state + t_interp * (target_state - initial_state)
initial_controls = jnp.zeros((n_steps, 1))

initial_vals = jaxls.VarValues.make(
    [
        state_vars.with_value(initial_states),
        control_vars.with_value(initial_controls),
    ]
)

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

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

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

INFO     | Building optimization problem with 252 terms and 121 variables: 72 costs, 60 eq_zero, 120 leq_zero, 0 geq_zero
INFO     | Vectorizing constraint group with 60 constraints (constraint_leq_zero), 1 variables each: augmented_control_upper_bound
INFO     | Vectorizing constraint group with 60 constraints (constraint_eq_zero), 3 variables each: augmented_dynamics_constraint
INFO     | Vectorizing group with 12 costs, 1 variables each: boundary_cost
INFO     | Vectorizing group with 60 costs, 1 variables each: control_cost
INFO     | Vectorizing constraint group with 60 constraints (constraint_leq_zero), 1 variables each: augmented_control_lower_bound
INFO     | Variable elimination: eliminating ControlVar (60 of 304 tangent dims); reduced system is 244-dimensional

Solving#

solution = problem.solve(initial_vals)

INFO     | Augmented Lagrangian: initial snorm=7.3573e-01, csupn=7.3573e-01, max_rho=1.0572e+05, constraint_dim=360
INFO     |  step #0: cost=10554.9863 lambd=0.0005 inexact_tol=1.0e-02
INFO     |      - augmented_control_upper_bound(60): 0.00000 (avg 0.00000)
INFO     |      - augmented_dynamics_constraint(60): 1805924.62500 (avg 7524.68652)
INFO     |      - boundary_cost(12): 10554.98633 (avg 219.89555)
INFO     |      - control_cost(60): 0.00000 (avg 0.00000)
INFO     |      - augmented_control_lower_bound(60): 0.00000 (avg 0.00000)
INFO     |  step #1: cost=10554.9863 lambd=0.0020 inexact_tol=1.0e-02
INFO     |      - augmented_control_upper_bound(60): 0.00000 (avg 0.00000)
INFO     |      - augmented_dynamics_constraint(60): 1805924.62500 (avg 7524.68652)
INFO     |      - boundary_cost(12): 10554.98633 (avg 219.89555)
INFO     |      - control_cost(60): 0.00000 (avg 0.00000)
INFO     |      - augmented_control_lower_bound(60): 0.00000 (avg 0.00000)
INFO     |  step #2: cost=10554.9863 lambd=0.0160 inexact_tol=1.0e-02
INFO     |      - augmented_control_upper_bound(60): 0.00000 (avg 0.00000)
INFO     |      - augmented_dynamics_constraint(60): 1805924.62500 (avg 7524.68652)
INFO     |      - boundary_cost(12): 10554.98633 (avg 219.89555)
INFO     |      - control_cost(60): 0.00000 (avg 0.00000)
INFO     |      - augmented_control_lower_bound(60): 0.00000 (avg 0.00000)
INFO     |  step #3: cost=10554.9863 lambd=0.2560 inexact_tol=1.0e-02
INFO     |      - augmented_control_upper_bound(60): 0.00000 (avg 0.00000)
INFO     |      - augmented_dynamics_constraint(60): 1805924.62500 (avg 7524.68652)
INFO     |      - boundary_cost(12): 10554.98633 (avg 219.89555)
INFO     |      - control_cost(60): 0.00000 (avg 0.00000)
INFO     |      - augmented_control_lower_bound(60): 0.00000 (avg 0.00000)
INFO     |  step #4: cost=10554.9863 lambd=8.1920 inexact_tol=1.0e-02
INFO     |      - augmented_control_upper_bound(60): 0.00000 (avg 0.00000)
INFO     |      - augmented_dynamics_constraint(60): 1805924.62500 (avg 7524.68652)
INFO     |      - boundary_cost(12): 10554.98633 (avg 219.89555)
INFO     |      - control_cost(60): 0.00000 (avg 0.00000)
INFO     |      - augmented_control_lower_bound(60): 0.00000 (avg 0.00000)
INFO     |      accepted=True ATb_norm=1.64e+05 cost_prev=1816479.6250 cost_new=1554753.1250
INFO     |  step #5: cost=3379.7109 lambd=4.0960 inexact_tol=1.0e-02
INFO     |      - augmented_control_upper_bound(60): 7775.60840 (avg 129.59348)
INFO     |      - augmented_dynamics_constraint(60): 1366023.50000 (avg 5691.76465)
INFO     |      - boundary_cost(12): 1439.42004 (avg 29.98792)
INFO     |      - control_cost(60): 1940.29102 (avg 32.33818)
INFO     |      - augmented_control_lower_bound(60): 177574.34375 (avg 2959.57251)
INFO     |      accepted=True ATb_norm=3.30e+05 cost_prev=1554753.1250 cost_new=61417.3438
INFO     |  step #6: cost=2495.5261 lambd=2.0480 inexact_tol=1.0e-02
INFO     |      - augmented_control_upper_bound(60): 35121.83594 (avg 585.36395)
INFO     |      - augmented_dynamics_constraint(60): 23799.98438 (avg 99.16660)
INFO     |      - boundary_cost(12): 640.12482 (avg 13.33593)
INFO     |      - control_cost(60): 1855.40137 (avg 30.92336)
INFO     |      - augmented_control_lower_bound(60): 0.00000 (avg 0.00000)
INFO     |      accepted=True ATb_norm=7.31e+04 cost_prev=61417.3438 cost_new=11155.5186
INFO     |  step #7: cost=1378.3356 lambd=1.0240 inexact_tol=1.0e-02
INFO     |      - augmented_control_upper_bound(60): 0.00000 (avg 0.00000)
INFO     |      - augmented_dynamics_constraint(60): 9777.18262 (avg 40.73826)
INFO     |      - boundary_cost(12): 104.89658 (avg 2.18535)
INFO     |      - control_cost(60): 1273.43896 (avg 21.22398)
INFO     |      - augmented_control_lower_bound(60): 0.00000 (avg 0.00000)
INFO     |  step #8: cost=1378.3356 lambd=4.0960 inexact_tol=1.0e-02
INFO     |      - augmented_control_upper_bound(60): 0.00000 (avg 0.00000)
INFO     |      - augmented_dynamics_constraint(60): 9777.18262 (avg 40.73826)
INFO     |      - boundary_cost(12): 104.89658 (avg 2.18535)
INFO     |      - control_cost(60): 1273.43896 (avg 21.22398)
INFO     |      - augmented_control_lower_bound(60): 0.00000 (avg 0.00000)
INFO     |      accepted=True ATb_norm=2.07e+04 cost_prev=11155.5186 cost_new=6117.8833
INFO     |  step #9: cost=1269.6632 lambd=2.0480 inexact_tol=1.0e-02
INFO     |      - augmented_control_upper_bound(60): 0.00000 (avg 0.00000)
INFO     |      - augmented_dynamics_constraint(60): 150.76300 (avg 0.62818)
INFO     |      - boundary_cost(12): 32.98902 (avg 0.68727)
INFO     |      - control_cost(60): 1236.67419 (avg 20.61124)
INFO     |      - augmented_control_lower_bound(60): 4697.45703 (avg 78.29095)
INFO     |  step #10: cost=1269.6632 lambd=8.1920 inexact_tol=1.0e-02
INFO     |      - augmented_control_upper_bound(60): 0.00000 (avg 0.00000)
INFO     |      - augmented_dynamics_constraint(60): 150.76300 (avg 0.62818)
INFO     |      - boundary_cost(12): 32.98902 (avg 0.68727)
INFO     |      - control_cost(60): 1236.67419 (avg 20.61124)
INFO     |      - augmented_control_lower_bound(60): 4697.45703 (avg 78.29095)
INFO     |      accepted=True ATb_norm=2.34e+04 cost_prev=6117.8833 cost_new=1298.6483
INFO     |  step #11: cost=1235.2429 lambd=4.0960 inexact_tol=1.0e-02
INFO     |      - augmented_control_upper_bound(60): 0.00000 (avg 0.00000)
INFO     |      - augmented_dynamics_constraint(60): 63.40487 (avg 0.26419)
INFO     |      - boundary_cost(12): 26.78951 (avg 0.55811)
INFO     |      - control_cost(60): 1208.45337 (avg 20.14089)
INFO     |      - augmented_control_lower_bound(60): 0.00045 (avg 0.00001)
INFO     |  step #12: cost=1235.2429 lambd=16.3840 inexact_tol=1.0e-02
INFO     |      - augmented_control_upper_bound(60): 0.00000 (avg 0.00000)
INFO     |      - augmented_dynamics_constraint(60): 63.40487 (avg 0.26419)
INFO     |      - boundary_cost(12): 26.78951 (avg 0.55811)
INFO     |      - control_cost(60): 1208.45337 (avg 20.14089)
INFO     |      - augmented_control_lower_bound(60): 0.00045 (avg 0.00001)
INFO     |      accepted=True ATb_norm=9.43e+01 cost_prev=1298.6483 cost_new=1276.8281
INFO     |  step #13: cost=1220.3676 lambd=8.1920 inexact_tol=2.0e-03
INFO     |      - augmented_control_upper_bound(60): 0.00000 (avg 0.00000)
INFO     |      - augmented_dynamics_constraint(60): 56.46043 (avg 0.23525)
INFO     |      - boundary_cost(12): 24.19597 (avg 0.50408)
INFO     |      - control_cost(60): 1196.17163 (avg 19.93620)
INFO     |      - augmented_control_lower_bound(60): 0.00019 (avg 0.00000)
INFO     |  step #14: cost=1220.3676 lambd=32.7680 inexact_tol=2.0e-03
INFO     |      - augmented_control_upper_bound(60): 0.00000 (avg 0.00000)
INFO     |      - augmented_dynamics_constraint(60): 56.46043 (avg 0.23525)
INFO     |      - boundary_cost(12): 24.19597 (avg 0.50408)
INFO     |      - control_cost(60): 1196.17163 (avg 19.93620)
INFO     |      - augmented_control_lower_bound(60): 0.00019 (avg 0.00000)
INFO     |      accepted=True ATb_norm=3.37e+01 cost_prev=1276.8281 cost_new=1267.3530
INFO     |  step #15: cost=1213.1959 lambd=16.3840 inexact_tol=2.0e-03
INFO     |      - augmented_control_upper_bound(60): 0.00000 (avg 0.00000)
INFO     |      - augmented_dynamics_constraint(60): 54.15703 (avg 0.22565)
INFO     |      - boundary_cost(12): 23.36483 (avg 0.48677)
INFO     |      - control_cost(60): 1189.83105 (avg 19.83052)
INFO     |      - augmented_control_lower_bound(60): 0.00016 (avg 0.00000)
INFO     |  step #16: cost=1213.1959 lambd=65.5360 inexact_tol=2.0e-03
INFO     |      - augmented_control_upper_bound(60): 0.00000 (avg 0.00000)
INFO     |      - augmented_dynamics_constraint(60): 54.15703 (avg 0.22565)
INFO     |      - boundary_cost(12): 23.36483 (avg 0.48677)
INFO     |      - control_cost(60): 1189.83105 (avg 19.83052)
INFO     |      - augmented_control_lower_bound(60): 0.00016 (avg 0.00000)
INFO     |  step #17: cost=1213.1959 lambd=524.2880 inexact_tol=2.0e-03
INFO     |      - augmented_control_upper_bound(60): 0.00000 (avg 0.00000)
INFO     |      - augmented_dynamics_constraint(60): 54.15703 (avg 0.22565)
INFO     |      - boundary_cost(12): 23.36483 (avg 0.48677)
INFO     |      - control_cost(60): 1189.83105 (avg 19.83052)
INFO     |      - augmented_control_lower_bound(60): 0.00016 (avg 0.00000)
INFO     |      accepted=True ATb_norm=1.39e+01 cost_prev=1267.3530 cost_new=1266.7814
INFO     |  step #18: cost=1212.7729 lambd=262.1440 inexact_tol=2.0e-03
INFO     |      - augmented_control_upper_bound(60): 0.00000 (avg 0.00000)
INFO     |      - augmented_dynamics_constraint(60): 54.00816 (avg 0.22503)
INFO     |      - boundary_cost(12): 23.33038 (avg 0.48605)
INFO     |      - control_cost(60): 1189.44263 (avg 19.82405)
INFO     |      - augmented_control_lower_bound(60): 0.00015 (avg 0.00000)
INFO     |      accepted=True ATb_norm=1.22e+01 cost_prev=1266.7814 cost_new=1265.6506
INFO     |  step #19: cost=1211.9062 lambd=131.0720 inexact_tol=2.0e-03
INFO     |      - augmented_control_upper_bound(60): 0.00000 (avg 0.00000)
INFO     |      - augmented_dynamics_constraint(60): 53.74420 (avg 0.22393)
INFO     |      - boundary_cost(12): 23.24915 (avg 0.48436)
INFO     |      - control_cost(60): 1188.65710 (avg 19.81095)
INFO     |      - augmented_control_lower_bound(60): 0.00015 (avg 0.00000)
INFO     |  step #20: cost=1211.9062 lambd=524.2880 inexact_tol=2.0e-03
INFO     |      - augmented_control_upper_bound(60): 0.00000 (avg 0.00000)
INFO     |      - augmented_dynamics_constraint(60): 53.74420 (avg 0.22393)
INFO     |      - boundary_cost(12): 23.24915 (avg 0.48436)
INFO     |      - control_cost(60): 1188.65710 (avg 19.81095)
INFO     |      - augmented_control_lower_bound(60): 0.00015 (avg 0.00000)
INFO     |      accepted=True ATb_norm=1.22e+01 cost_prev=1265.6506 cost_new=1265.6238
INFO     |  step #21: cost=1211.4728 lambd=262.1440 inexact_tol=2.0e-03
INFO     |      - augmented_control_upper_bound(60): 0.00000 (avg 0.00000)
INFO     |      - augmented_dynamics_constraint(60): 53.61704 (avg 0.22340)
INFO     |      - boundary_cost(12): 23.20944 (avg 0.48353)
INFO     |      - control_cost(60): 1188.26331 (avg 19.80439)
INFO     |      - augmented_control_lower_bound(60): 0.53386 (avg 0.00890)
INFO     |      accepted=True ATb_norm=2.47e+02 cost_prev=1265.6238 cost_new=1264.0225
INFO     |  step #22: cost=1210.4347 lambd=131.0720 inexact_tol=2.0e-03
INFO     |      - augmented_control_upper_bound(60): 0.00000 (avg 0.00000)
INFO     |      - augmented_dynamics_constraint(60): 53.58755 (avg 0.22328)
INFO     |      - boundary_cost(12): 23.15535 (avg 0.48240)
INFO     |      - control_cost(60): 1187.27930 (avg 19.78799)
INFO     |      - augmented_control_lower_bound(60): 0.00022 (avg 0.00000)
INFO     |      accepted=True ATb_norm=1.19e+01 cost_prev=1264.0225 cost_new=1261.9067
INFO     |  step #23: cost=1208.4780 lambd=65.5360 inexact_tol=2.0e-03
INFO     |      - augmented_control_upper_bound(60): 0.00000 (avg 0.00000)
INFO     |      - augmented_dynamics_constraint(60): 53.42867 (avg 0.22262)
INFO     |      - boundary_cost(12): 23.05958 (avg 0.48041)
INFO     |      - control_cost(60): 1185.41846 (avg 19.75698)
INFO     |      - augmented_control_lower_bound(60): 0.00019 (avg 0.00000)
INFO     |      accepted=True ATb_norm=1.18e+01 cost_prev=1261.9067 cost_new=1257.8201
INFO     |  step #24: cost=1204.6868 lambd=32.7680 inexact_tol=2.0e-03
INFO     |      - augmented_control_upper_bound(60): 0.00000 (avg 0.00000)
INFO     |      - augmented_dynamics_constraint(60): 53.13321 (avg 0.22139)
INFO     |      - boundary_cost(12): 22.89448 (avg 0.47697)
INFO     |      - control_cost(60): 1181.79224 (avg 19.69654)
INFO     |      - augmented_control_lower_bound(60): 0.00020 (avg 0.00000)
INFO     |      accepted=True ATb_norm=1.16e+01 cost_prev=1257.8201 cost_new=1250.1804
INFO     |  step #25: cost=1197.5540 lambd=16.3840 inexact_tol=2.0e-03
INFO     |      - augmented_control_upper_bound(60): 0.00000 (avg 0.00000)
INFO     |      - augmented_dynamics_constraint(60): 52.62607 (avg 0.21928)
INFO     |      - boundary_cost(12): 22.64056 (avg 0.47168)
INFO     |      - control_cost(60): 1174.91345 (avg 19.58189)
INFO     |      - augmented_control_lower_bound(60): 0.00020 (avg 0.00000)
INFO     |      accepted=True ATb_norm=1.20e+01 cost_prev=1250.1804 cost_new=1236.7474
INFO     |  step #26: cost=1184.8229 lambd=8.1920 inexact_tol=2.0e-03
INFO     |      - augmented_control_upper_bound(60): 0.00000 (avg 0.00000)
INFO     |      - augmented_dynamics_constraint(60): 51.92427 (avg 0.21635)
INFO     |      - boundary_cost(12): 22.30600 (avg 0.46471)
INFO     |      - control_cost(60): 1162.51685 (avg 19.37528)
INFO     |      - augmented_control_lower_bound(60): 0.00020 (avg 0.00000)
INFO     |      accepted=True ATb_norm=1.73e+01 cost_prev=1236.7474 cost_new=1215.6492
INFO     |  step #27: cost=1164.3118 lambd=4.0960 inexact_tol=2.0e-03
INFO     |      - augmented_control_upper_bound(60): 0.00000 (avg 0.00000)
INFO     |      - augmented_dynamics_constraint(60): 51.33718 (avg 0.21390)
INFO     |      - boundary_cost(12): 22.01349 (avg 0.45861)
INFO     |      - control_cost(60): 1142.29834 (avg 19.03831)
INFO     |      - augmented_control_lower_bound(60): 0.00019 (avg 0.00000)
INFO     |      accepted=True ATb_norm=3.86e+01 cost_prev=1215.6492 cost_new=1188.5734
INFO     |  step #28: cost=1137.1696 lambd=2.0480 inexact_tol=2.0e-03
INFO     |      - augmented_control_upper_bound(60): 0.00000 (avg 0.00000)
INFO     |      - augmented_dynamics_constraint(60): 51.40349 (avg 0.21418)
INFO     |      - boundary_cost(12): 22.00821 (avg 0.45850)
INFO     |      - control_cost(60): 1115.16138 (avg 18.58602)
INFO     |      - augmented_control_lower_bound(60): 0.00019 (avg 0.00000)
INFO     |      accepted=True ATb_norm=8.26e+01 cost_prev=1188.5734 cost_new=1163.7723
INFO     |  step #29: cost=1111.6046 lambd=1.0240 inexact_tol=2.0e-03
INFO     |      - augmented_control_upper_bound(60): 0.00000 (avg 0.00000)
INFO     |      - augmented_dynamics_constraint(60): 52.16747 (avg 0.21736)
INFO     |      - boundary_cost(12): 22.41615 (avg 0.46700)
INFO     |      - control_cost(60): 1089.18848 (avg 18.15314)
INFO     |      - augmented_control_lower_bound(60): 0.00019 (avg 0.00000)
INFO     |  step #30: cost=1111.6046 lambd=4.0960 inexact_tol=2.0e-03
INFO     |      - augmented_control_upper_bound(60): 0.00000 (avg 0.00000)
INFO     |      - augmented_dynamics_constraint(60): 52.16747 (avg 0.21736)
INFO     |      - boundary_cost(12): 22.41615 (avg 0.46700)
INFO     |      - control_cost(60): 1089.18848 (avg 18.15314)
INFO     |      - augmented_control_lower_bound(60): 0.00019 (avg 0.00000)
INFO     |      accepted=True ATb_norm=1.27e+02 cost_prev=1163.7723 cost_new=1156.5872
INFO     |  step #31: cost=1104.9604 lambd=2.0480 inexact_tol=2.0e-03
INFO     |      - augmented_control_upper_bound(60): 0.00000 (avg 0.00000)
INFO     |      - augmented_dynamics_constraint(60): 51.62655 (avg 0.21511)
INFO     |      - boundary_cost(12): 22.83932 (avg 0.47582)
INFO     |      - control_cost(60): 1082.12109 (avg 18.03535)
INFO     |      - augmented_control_lower_bound(60): 0.00018 (avg 0.00000)
INFO     |  step #32: cost=1104.9604 lambd=8.1920 inexact_tol=2.0e-03
INFO     |      - augmented_control_upper_bound(60): 0.00000 (avg 0.00000)
INFO     |      - augmented_dynamics_constraint(60): 51.62655 (avg 0.21511)
INFO     |      - boundary_cost(12): 22.83932 (avg 0.47582)
INFO     |      - control_cost(60): 1082.12109 (avg 18.03535)
INFO     |      - augmented_control_lower_bound(60): 0.00018 (avg 0.00000)
INFO     |      accepted=True ATb_norm=1.87e+01 cost_prev=1156.5872 cost_new=1154.1467
INFO     |  step #33: cost=1102.4559 lambd=4.0960 inexact_tol=2.0e-03
INFO     |      - augmented_control_upper_bound(60): 0.00000 (avg 0.00000)
INFO     |      - augmented_dynamics_constraint(60): 51.69062 (avg 0.21538)
INFO     |      - boundary_cost(12): 22.88637 (avg 0.47680)
INFO     |      - control_cost(60): 1079.56958 (avg 17.99283)
INFO     |      - augmented_control_lower_bound(60): 0.00018 (avg 0.00000)
INFO     |  step #34: cost=1102.4559 lambd=16.3840 inexact_tol=2.0e-03
INFO     |      - augmented_control_upper_bound(60): 0.00000 (avg 0.00000)
INFO     |      - augmented_dynamics_constraint(60): 51.69062 (avg 0.21538)
INFO     |      - boundary_cost(12): 22.88637 (avg 0.47680)
INFO     |      - control_cost(60): 1079.56958 (avg 17.99283)
INFO     |      - augmented_control_lower_bound(60): 0.00018 (avg 0.00000)
INFO     |  step #35: cost=1102.4559 lambd=131.0720 inexact_tol=2.0e-03
INFO     |      - augmented_control_upper_bound(60): 0.00000 (avg 0.00000)
INFO     |      - augmented_dynamics_constraint(60): 51.69062 (avg 0.21538)
INFO     |      - boundary_cost(12): 22.88637 (avg 0.47680)
INFO     |      - control_cost(60): 1079.56958 (avg 17.99283)
INFO     |      - augmented_control_lower_bound(60): 0.00018 (avg 0.00000)
INFO     |      accepted=True ATb_norm=5.01e+00 cost_prev=1154.1467 cost_new=1154.0066
INFO     |  step #36: cost=1102.3107 lambd=65.5360 inexact_tol=2.0e-03
INFO     |      - augmented_control_upper_bound(60): 0.00000 (avg 0.00000)
INFO     |      - augmented_dynamics_constraint(60): 51.69563 (avg 0.21540)
INFO     |      - boundary_cost(12): 22.88481 (avg 0.47677)
INFO     |      - control_cost(60): 1079.42590 (avg 17.99043)
INFO     |      - augmented_control_lower_bound(60): 0.00018 (avg 0.00000)
INFO     |      accepted=True ATb_norm=3.04e+00 cost_prev=1154.0066 cost_new=1153.7822
INFO     |  step #37: cost=1102.0280 lambd=32.7680 inexact_tol=2.0e-03
INFO     |      - augmented_control_upper_bound(60): 0.00000 (avg 0.00000)
INFO     |      - augmented_dynamics_constraint(60): 51.70588 (avg 0.21544)
INFO     |      - boundary_cost(12): 22.88643 (avg 0.47680)
INFO     |      - control_cost(60): 1079.14148 (avg 17.98569)
INFO     |      - augmented_control_lower_bound(60): 0.04837 (avg 0.00081)
INFO     |      accepted=True ATb_norm=7.42e+01 cost_prev=1153.7822 cost_new=1153.2375
INFO     |  step #38: cost=1101.2941 lambd=16.3840 inexact_tol=2.0e-03
INFO     |      - augmented_control_upper_bound(60): 0.00000 (avg 0.00000)
INFO     |      - augmented_dynamics_constraint(60): 51.94338 (avg 0.21643)
INFO     |      - boundary_cost(12): 22.95756 (avg 0.47828)
INFO     |      - control_cost(60): 1078.33655 (avg 17.97228)
INFO     |      - augmented_control_lower_bound(60): 0.00019 (avg 0.00000)
INFO     |      accepted=True ATb_norm=2.87e+00 cost_prev=1153.2375 cost_new=1152.3481
INFO     |  step #39: cost=1100.1207 lambd=8.1920 inexact_tol=2.0e-03
INFO     |      - augmented_control_upper_bound(60): 0.00000 (avg 0.00000)
INFO     |      - augmented_dynamics_constraint(60): 52.22722 (avg 0.21761)
INFO     |      - boundary_cost(12): 23.06812 (avg 0.48059)
INFO     |      - control_cost(60): 1077.05261 (avg 17.95088)
INFO     |      - augmented_control_lower_bound(60): 0.00019 (avg 0.00000)
INFO     |      accepted=True ATb_norm=2.83e+00 cost_prev=1152.3481 cost_new=1150.9125
INFO     |  step #40: cost=1098.3901 lambd=4.0960 inexact_tol=2.0e-03
INFO     |      - augmented_control_upper_bound(60): 0.00000 (avg 0.00000)
INFO     |      - augmented_dynamics_constraint(60): 52.52209 (avg 0.21884)
INFO     |      - boundary_cost(12): 23.19939 (avg 0.48332)
INFO     |      - control_cost(60): 1075.19080 (avg 17.91985)
INFO     |      - augmented_control_lower_bound(60): 0.00020 (avg 0.00000)
INFO     |      accepted=True ATb_norm=3.42e+00 cost_prev=1150.9125 cost_new=1148.9894
INFO     |  step #41: cost=1096.2053 lambd=2.0480 inexact_tol=2.0e-03
INFO     |      - augmented_control_upper_bound(60): 0.00000 (avg 0.00000)
INFO     |      - augmented_dynamics_constraint(60): 52.78383 (avg 0.21993)
INFO     |      - boundary_cost(12): 23.30614 (avg 0.48554)
INFO     |      - control_cost(60): 1072.89917 (avg 17.88165)
INFO     |      - augmented_control_lower_bound(60): 0.00020 (avg 0.00000)
INFO     |      accepted=True ATb_norm=6.27e+00 cost_prev=1148.9894 cost_new=1147.1204
INFO     |  step #42: cost=1094.1281 lambd=1.0240 inexact_tol=2.0e-03
INFO     |      - augmented_control_upper_bound(60): 0.00000 (avg 0.00000)
INFO     |      - augmented_dynamics_constraint(60): 52.99198 (avg 0.22080)
INFO     |      - boundary_cost(12): 23.38408 (avg 0.48717)
INFO     |      - control_cost(60): 1070.74402 (avg 17.84573)
INFO     |      - augmented_control_lower_bound(60): 0.00021 (avg 0.00000)
INFO     |      accepted=True ATb_norm=1.09e+01 cost_prev=1147.1204 cost_new=1146.0040
INFO     |  step #43: cost=1092.8716 lambd=0.5120 inexact_tol=2.0e-03
INFO     |      - augmented_control_upper_bound(60): 0.00000 (avg 0.00000)
INFO     |      - augmented_dynamics_constraint(60): 53.13224 (avg 0.22138)
INFO     |      - boundary_cost(12): 23.43310 (avg 0.48819)
INFO     |      - control_cost(60): 1069.43848 (avg 17.82397)
INFO     |      - augmented_control_lower_bound(60): 0.00021 (avg 0.00000)
INFO     |      accepted=True ATb_norm=1.23e+01 cost_prev=1146.0040 cost_new=1145.6370
INFO     |  step #44: cost=1092.4552 lambd=0.2560 inexact_tol=2.0e-03
INFO     |      - augmented_control_upper_bound(60): 0.00000 (avg 0.00000)
INFO     |      - augmented_dynamics_constraint(60): 53.18153 (avg 0.22159)
INFO     |      - boundary_cost(12): 23.43345 (avg 0.48820)
INFO     |      - control_cost(60): 1069.02173 (avg 17.81703)
INFO     |      - augmented_control_lower_bound(60): 0.00020 (avg 0.00000)
INFO     |      accepted=True ATb_norm=7.92e+00 cost_prev=1145.6370 cost_new=1145.5426
INFO     |  step #45: cost=1092.3754 lambd=0.1280 inexact_tol=2.0e-03
INFO     |      - augmented_control_upper_bound(60): 0.00000 (avg 0.00000)
INFO     |      - augmented_dynamics_constraint(60): 53.16706 (avg 0.22153)
INFO     |      - boundary_cost(12): 23.41233 (avg 0.48776)
INFO     |      - control_cost(60): 1068.96301 (avg 17.81605)
INFO     |      - augmented_control_lower_bound(60): 0.00020 (avg 0.00000)
INFO     |      accepted=True ATb_norm=4.07e+00 cost_prev=1145.5426 cost_new=1145.4890
INFO     |  step #46: cost=1092.3290 lambd=0.0640 inexact_tol=2.0e-03
INFO     |      - augmented_control_upper_bound(60): 0.00000 (avg 0.00000)
INFO     |      - augmented_dynamics_constraint(60): 53.15985 (avg 0.22150)
INFO     |      - boundary_cost(12): 23.45109 (avg 0.48856)
INFO     |      - control_cost(60): 1068.87793 (avg 17.81463)
INFO     |      - augmented_control_lower_bound(60): 0.00020 (avg 0.00000)
INFO     |      accepted=True ATb_norm=3.42e+00 cost_prev=1145.4890 cost_new=1145.4401
INFO     |  step #47: cost=1092.2655 lambd=0.0320 inexact_tol=2.0e-03
INFO     |      - augmented_control_upper_bound(60): 0.00000 (avg 0.00000)
INFO     |      - augmented_dynamics_constraint(60): 53.17432 (avg 0.22156)
INFO     |      - boundary_cost(12): 23.54459 (avg 0.49051)
INFO     |      - control_cost(60): 1068.72095 (avg 17.81202)
INFO     |      - augmented_control_lower_bound(60): 0.00020 (avg 0.00000)
INFO     |      accepted=True ATb_norm=3.30e+00 cost_prev=1145.4401 cost_new=1145.3953
INFO     |  step #48: cost=1092.2026 lambd=0.0160 inexact_tol=2.0e-03
INFO     |      - augmented_control_upper_bound(60): 0.00000 (avg 0.00000)
INFO     |      - augmented_dynamics_constraint(60): 53.19245 (avg 0.22164)
INFO     |      - boundary_cost(12): 23.64884 (avg 0.49268)
INFO     |      - control_cost(60): 1068.55383 (avg 17.80923)
INFO     |      - augmented_control_lower_bound(60): 0.00020 (avg 0.00000)
INFO     |      accepted=True ATb_norm=3.15e+00 cost_prev=1145.3953 cost_new=1145.3545
INFO     |  step #49: cost=1092.1448 lambd=0.0080 inexact_tol=2.0e-03
INFO     |      - augmented_control_upper_bound(60): 0.00000 (avg 0.00000)
INFO     |      - augmented_dynamics_constraint(60): 53.20946 (avg 0.22171)
INFO     |      - boundary_cost(12): 23.75200 (avg 0.49483)
INFO     |      - control_cost(60): 1068.39282 (avg 17.80655)
INFO     |      - augmented_control_lower_bound(60): 0.00020 (avg 0.00000)
INFO     |      accepted=True ATb_norm=2.91e+00 cost_prev=1145.3545 cost_new=1145.3181
INFO     |  step #50: cost=1092.0939 lambd=0.0040 inexact_tol=2.0e-03
INFO     |      - augmented_control_upper_bound(60): 0.00000 (avg 0.00000)
INFO     |      - augmented_dynamics_constraint(60): 53.22404 (avg 0.22177)
INFO     |      - boundary_cost(12): 23.84886 (avg 0.49685)
INFO     |      - control_cost(60): 1068.24500 (avg 17.80408)
INFO     |      - augmented_control_lower_bound(60): 0.00020 (avg 0.00000)
INFO     |      accepted=True ATb_norm=2.63e+00 cost_prev=1145.3181 cost_new=1145.2858
INFO     |  step #51: cost=1092.0494 lambd=0.0020 inexact_tol=2.0e-03
INFO     |      - augmented_control_upper_bound(60): 0.00000 (avg 0.00000)
INFO     |      - augmented_dynamics_constraint(60): 53.23621 (avg 0.22182)
INFO     |      - boundary_cost(12): 23.93983 (avg 0.49875)
INFO     |      - control_cost(60): 1068.10962 (avg 17.80183)
INFO     |      - augmented_control_lower_bound(60): 0.00020 (avg 0.00000)
INFO     |  step #52: cost=1092.0494 lambd=0.0080 inexact_tol=2.0e-03
INFO     |      - augmented_control_upper_bound(60): 0.00000 (avg 0.00000)
INFO     |      - augmented_dynamics_constraint(60): 53.23621 (avg 0.22182)
INFO     |      - boundary_cost(12): 23.93983 (avg 0.49875)
INFO     |      - control_cost(60): 1068.10962 (avg 17.80183)
INFO     |      - augmented_control_lower_bound(60): 0.00020 (avg 0.00000)
INFO     |  step #53: cost=1092.0494 lambd=0.0640 inexact_tol=2.0e-03
INFO     |      - augmented_control_upper_bound(60): 0.00000 (avg 0.00000)
INFO     |      - augmented_dynamics_constraint(60): 53.23621 (avg 0.22182)
INFO     |      - boundary_cost(12): 23.93983 (avg 0.49875)
INFO     |      - control_cost(60): 1068.10962 (avg 17.80183)
INFO     |      - augmented_control_lower_bound(60): 0.00020 (avg 0.00000)
INFO     |  step #54: cost=1092.0494 lambd=1.0240 inexact_tol=2.0e-03
INFO     |      - augmented_control_upper_bound(60): 0.00000 (avg 0.00000)
INFO     |      - augmented_dynamics_constraint(60): 53.23621 (avg 0.22182)
INFO     |      - boundary_cost(12): 23.93983 (avg 0.49875)
INFO     |      - control_cost(60): 1068.10962 (avg 17.80183)
INFO     |      - augmented_control_lower_bound(60): 0.00020 (avg 0.00000)
INFO     |      accepted=True ATb_norm=2.32e+00 cost_prev=1145.2858 cost_new=1145.2710
INFO     |  step #55: cost=1092.0347 lambd=0.5120 inexact_tol=2.0e-03
INFO     |      - augmented_control_upper_bound(60): 0.00000 (avg 0.00000)
INFO     |      - augmented_dynamics_constraint(60): 53.23617 (avg 0.22182)
INFO     |      - boundary_cost(12): 23.99957 (avg 0.49999)
INFO     |      - control_cost(60): 1068.03516 (avg 17.80059)
INFO     |      - augmented_control_lower_bound(60): 0.00019 (avg 0.00000)
INFO     |  step #56: cost=1092.0347 lambd=2.0480 inexact_tol=2.0e-03
INFO     |      - augmented_control_upper_bound(60): 0.00000 (avg 0.00000)
INFO     |      - augmented_dynamics_constraint(60): 53.23617 (avg 0.22182)
INFO     |      - boundary_cost(12): 23.99957 (avg 0.49999)
INFO     |      - control_cost(60): 1068.03516 (avg 17.80059)
INFO     |      - augmented_control_lower_bound(60): 0.00019 (avg 0.00000)
INFO     |      accepted=True ATb_norm=7.35e-01 cost_prev=1145.2710 cost_new=1145.2622
INFO     |  step #57: cost=1092.0234 lambd=1.0240 inexact_tol=2.0e-03
INFO     |      - augmented_control_upper_bound(60): 0.00000 (avg 0.00000)
INFO     |      - augmented_dynamics_constraint(60): 53.23866 (avg 0.22183)
INFO     |      - boundary_cost(12): 24.03544 (avg 0.50074)
INFO     |      - control_cost(60): 1067.98804 (avg 17.79980)
INFO     |      - augmented_control_lower_bound(60): 0.00019 (avg 0.00000)
INFO     |  step #58: cost=1092.0234 lambd=4.0960 inexact_tol=2.0e-03
INFO     |      - augmented_control_upper_bound(60): 0.00000 (avg 0.00000)
INFO     |      - augmented_dynamics_constraint(60): 53.23866 (avg 0.22183)
INFO     |      - boundary_cost(12): 24.03544 (avg 0.50074)
INFO     |      - control_cost(60): 1067.98804 (avg 17.79980)
INFO     |      - augmented_control_lower_bound(60): 0.00019 (avg 0.00000)
INFO     |  step #59: cost=1092.0234 lambd=32.7680 inexact_tol=2.0e-03
INFO     |      - augmented_control_upper_bound(60): 0.00000 (avg 0.00000)
INFO     |      - augmented_dynamics_constraint(60): 53.23866 (avg 0.22183)
INFO     |      - boundary_cost(12): 24.03544 (avg 0.50074)
INFO     |      - control_cost(60): 1067.98804 (avg 17.79980)
INFO     |      - augmented_control_lower_bound(60): 0.00019 (avg 0.00000)
INFO     |      accepted=True ATb_norm=5.12e-01 cost_prev=1145.2622 cost_new=1145.2615
INFO     |  AL update: snorm=5.1892e-03, csupn=5.1892e-03, max_rho=4.2288e+05
INFO     |  step #60: cost=1092.0265 lambd=16.3840 inexact_tol=2.0e-03
INFO     |      - augmented_control_upper_bound(60): 0.00000 (avg 0.00000)
INFO     |      - augmented_dynamics_constraint(60): 212.93916 (avg 0.88725)
INFO     |      - boundary_cost(12): 24.04065 (avg 0.50085)
INFO     |      - control_cost(60): 1067.98584 (avg 17.79976)
INFO     |      - augmented_control_lower_bound(60): 0.00116 (avg 0.00002)
INFO     |  step #61: cost=1092.0265 lambd=65.5360 inexact_tol=2.0e-03
INFO     |      - augmented_control_upper_bound(60): 0.00000 (avg 0.00000)
INFO     |      - augmented_dynamics_constraint(60): 212.93916 (avg 0.88725)
INFO     |      - boundary_cost(12): 24.04065 (avg 0.50085)
INFO     |      - control_cost(60): 1067.98584 (avg 17.79976)
INFO     |      - augmented_control_lower_bound(60): 0.00116 (avg 0.00002)
INFO     |  step #62: cost=1092.0265 lambd=524.2880 inexact_tol=2.0e-03
INFO     |      - augmented_control_upper_bound(60): 0.00000 (avg 0.00000)
INFO     |      - augmented_dynamics_constraint(60): 212.93916 (avg 0.88725)
INFO     |      - boundary_cost(12): 24.04065 (avg 0.50085)
INFO     |      - control_cost(60): 1067.98584 (avg 17.79976)
INFO     |      - augmented_control_lower_bound(60): 0.00116 (avg 0.00002)
INFO     |  step #63: cost=1092.0265 lambd=8388.6084 inexact_tol=2.0e-03
INFO     |      - augmented_control_upper_bound(60): 0.00000 (avg 0.00000)
INFO     |      - augmented_dynamics_constraint(60): 212.93916 (avg 0.88725)
INFO     |      - boundary_cost(12): 24.04065 (avg 0.50085)
INFO     |      - control_cost(60): 1067.98584 (avg 17.79976)
INFO     |      - augmented_control_lower_bound(60): 0.00116 (avg 0.00002)
INFO     |      accepted=True ATb_norm=4.92e+02 cost_prev=1304.9668 cost_new=1300.6080
INFO     |  step #64: cost=1099.0911 lambd=4194.3042 inexact_tol=2.0e-03
INFO     |      - augmented_control_upper_bound(60): 0.00000 (avg 0.00000)
INFO     |      - augmented_dynamics_constraint(60): 201.51514 (avg 0.83965)
INFO     |      - boundary_cost(12): 30.96261 (avg 0.64505)
INFO     |      - control_cost(60): 1068.12842 (avg 17.80214)
INFO     |      - augmented_control_lower_bound(60): 0.00183 (avg 0.00003)
INFO     |      accepted=True ATb_norm=9.60e+01 cost_prev=1300.6080 cost_new=1298.8513
INFO     |  step #65: cost=1103.0721 lambd=2097.1521 inexact_tol=2.0e-03
INFO     |      - augmented_control_upper_bound(60): 0.00000 (avg 0.00000)
INFO     |      - augmented_dynamics_constraint(60): 195.59070 (avg 0.81496)
INFO     |      - boundary_cost(12): 34.67389 (avg 0.72237)
INFO     |      - control_cost(60): 1068.39832 (avg 17.80664)
INFO     |      - augmented_control_lower_bound(60): 0.18836 (avg 0.00314)
INFO     |      accepted=True ATb_norm=1.49e+02 cost_prev=1298.8513 cost_new=1296.6476
INFO     |  step #66: cost=1106.8933 lambd=1048.5760 inexact_tol=2.0e-03
INFO     |      - augmented_control_upper_bound(60): 0.00000 (avg 0.00000)
INFO     |      - augmented_dynamics_constraint(60): 189.75114 (avg 0.79063)
INFO     |      - boundary_cost(12): 38.09868 (avg 0.79372)
INFO     |      - control_cost(60): 1068.79468 (avg 17.81325)
INFO     |      - augmented_control_lower_bound(60): 0.00305 (avg 0.00005)
INFO     |      accepted=True ATb_norm=4.36e+01 cost_prev=1296.6476 cost_new=1294.2656
INFO     |  step #67: cost=1111.5350 lambd=524.2880 inexact_tol=2.0e-03
INFO     |      - augmented_control_upper_bound(60): 0.00000 (avg 0.00000)
INFO     |      - augmented_dynamics_constraint(60): 182.72815 (avg 0.76137)
INFO     |      - boundary_cost(12): 41.92250 (avg 0.87339)
INFO     |      - control_cost(60): 1069.61255 (avg 17.82688)
INFO     |      - augmented_control_lower_bound(60): 0.00235 (avg 0.00004)
INFO     |      accepted=True ATb_norm=3.31e+01 cost_prev=1294.2656 cost_new=1291.5371
INFO     |  step #68: cost=1117.1023 lambd=262.1440 inexact_tol=2.0e-03
INFO     |      - augmented_control_upper_bound(60): 0.00000 (avg 0.00000)
INFO     |      - augmented_dynamics_constraint(60): 174.43298 (avg 0.72680)
INFO     |      - boundary_cost(12): 45.90231 (avg 0.95630)
INFO     |      - control_cost(60): 1071.19995 (avg 17.85333)
INFO     |      - augmented_control_lower_bound(60): 0.00206 (avg 0.00003)
INFO     |      accepted=True ATb_norm=2.51e+01 cost_prev=1291.5371 cost_new=1288.1023
INFO     |  step #69: cost=1123.4437 lambd=131.0720 inexact_tol=2.0e-03
INFO     |      - augmented_control_upper_bound(60): 0.00000 (avg 0.00000)
INFO     |      - augmented_dynamics_constraint(60): 164.65672 (avg 0.68607)
INFO     |      - boundary_cost(12): 49.27234 (avg 1.02651)
INFO     |      - control_cost(60): 1074.17139 (avg 17.90286)
INFO     |      - augmented_control_lower_bound(60): 0.00183 (avg 0.00003)
INFO     |      accepted=True ATb_norm=2.03e+01 cost_prev=1288.1023 cost_new=1283.4492
INFO     |  step #70: cost=1130.9563 lambd=65.5360 inexact_tol=2.0e-03
INFO     |      - augmented_control_upper_bound(60): 0.00000 (avg 0.00000)
INFO     |      - augmented_dynamics_constraint(60): 152.49135 (avg 0.63538)
INFO     |      - boundary_cost(12): 51.53373 (avg 1.07362)
INFO     |      - control_cost(60): 1079.42261 (avg 17.99038)
INFO     |      - augmented_control_lower_bound(60): 0.00154 (avg 0.00003)
INFO     |      accepted=True ATb_norm=1.67e+01 cost_prev=1283.4492 cost_new=1277.6361
INFO     |  step #71: cost=1140.3032 lambd=32.7680 inexact_tol=2.0e-03
INFO     |      - augmented_control_upper_bound(60): 0.00000 (avg 0.00000)
INFO     |      - augmented_dynamics_constraint(60): 137.33179 (avg 0.57222)
INFO     |      - boundary_cost(12): 52.22104 (avg 1.08794)
INFO     |      - control_cost(60): 1088.08215 (avg 18.13470)
INFO     |      - augmented_control_lower_bound(60): 0.00123 (avg 0.00002)
INFO     |      accepted=True ATb_norm=1.31e+01 cost_prev=1277.6361 cost_new=1271.1927
INFO     |  step #72: cost=1151.1964 lambd=16.3840 inexact_tol=2.0e-03
INFO     |      - augmented_control_upper_bound(60): 0.00000 (avg 0.00000)
INFO     |      - augmented_dynamics_constraint(60): 119.99563 (avg 0.49998)
INFO     |      - boundary_cost(12): 50.05529 (avg 1.04282)
INFO     |      - control_cost(60): 1101.14111 (avg 18.35235)
INFO     |      - augmented_control_lower_bound(60): 0.00085 (avg 0.00001)
INFO     |      accepted=True ATb_norm=1.07e+01 cost_prev=1271.1927 cost_new=1265.0737
INFO     |  step #73: cost=1162.8845 lambd=8.1920 inexact_tol=2.0e-03
INFO     |      - augmented_control_upper_bound(60): 0.00000 (avg 0.00000)
INFO     |      - augmented_dynamics_constraint(60): 102.18878 (avg 0.42579)
INFO     |      - boundary_cost(12): 44.48523 (avg 0.92678)
INFO     |      - control_cost(60): 1118.39929 (avg 18.63999)
INFO     |      - augmented_control_lower_bound(60): 0.00052 (avg 0.00001)
INFO     |      accepted=True ATb_norm=1.79e+01 cost_prev=1265.0737 cost_new=1260.8894
INFO     |  step #74: cost=1174.5490 lambd=4.0960 inexact_tol=2.0e-03
INFO     |      - augmented_control_upper_bound(60): 0.00000 (avg 0.00000)
INFO     |      - augmented_dynamics_constraint(60): 86.34004 (avg 0.35975)
INFO     |      - boundary_cost(12): 37.77410 (avg 0.78696)
INFO     |      - control_cost(60): 1136.77490 (avg 18.94625)
INFO     |      - augmented_control_lower_bound(60): 0.00028 (avg 0.00000)
INFO     |      accepted=True ATb_norm=2.29e+01 cost_prev=1260.8894 cost_new=1259.1748
INFO     |  step #75: cost=1183.7344 lambd=2.0480 inexact_tol=2.0e-03
INFO     |      - augmented_control_upper_bound(60): 0.00000 (avg 0.00000)
INFO     |      - augmented_dynamics_constraint(60): 75.44025 (avg 0.31433)
INFO     |      - boundary_cost(12): 32.93917 (avg 0.68623)
INFO     |      - control_cost(60): 1150.79517 (avg 19.17992)
INFO     |      - augmented_control_lower_bound(60): 0.00017 (avg 0.00000)
INFO     |      accepted=True ATb_norm=1.18e+01 cost_prev=1259.1748 cost_new=1258.7456
INFO     |  step #76: cost=1188.3903 lambd=1.0240 inexact_tol=2.0e-03
INFO     |      - augmented_control_upper_bound(60): 0.00000 (avg 0.00000)
INFO     |      - augmented_dynamics_constraint(60): 70.35531 (avg 0.29315)
INFO     |      - boundary_cost(12): 30.73210 (avg 0.64025)
INFO     |      - control_cost(60): 1157.65820 (avg 19.29430)
INFO     |      - augmented_control_lower_bound(60): 0.00013 (avg 0.00000)
INFO     |      accepted=True ATb_norm=3.11e+00 cost_prev=1258.7456 cost_new=1258.6313
INFO     |  step #77: cost=1189.8330 lambd=0.5120 inexact_tol=2.0e-03
INFO     |      - augmented_control_upper_bound(60): 0.00000 (avg 0.00000)
INFO     |      - augmented_dynamics_constraint(60): 68.79825 (avg 0.28666)
INFO     |      - boundary_cost(12): 29.94674 (avg 0.62389)
INFO     |      - control_cost(60): 1159.88623 (avg 19.33144)
INFO     |      - augmented_control_lower_bound(60): 0.00013 (avg 0.00000)
INFO     |      accepted=True ATb_norm=2.42e+00 cost_prev=1258.6313 cost_new=1258.5767
INFO     |  step #78: cost=1190.2522 lambd=0.2560 inexact_tol=2.0e-03
INFO     |      - augmented_control_upper_bound(60): 0.00000 (avg 0.00000)
INFO     |      - augmented_dynamics_constraint(60): 68.32446 (avg 0.28469)
INFO     |      - boundary_cost(12): 29.52794 (avg 0.61517)
INFO     |      - control_cost(60): 1160.72424 (avg 19.34541)
INFO     |      - augmented_control_lower_bound(60): 0.00013 (avg 0.00000)
INFO     |      accepted=True ATb_norm=3.63e+00 cost_prev=1258.5767 cost_new=1258.5291
INFO     |  step #79: cost=1190.4254 lambd=0.1280 inexact_tol=2.0e-03
INFO     |      - augmented_control_upper_bound(60): 0.00000 (avg 0.00000)
INFO     |      - augmented_dynamics_constraint(60): 68.10345 (avg 0.28376)
INFO     |      - boundary_cost(12): 29.27110 (avg 0.60981)
INFO     |      - control_cost(60): 1161.15430 (avg 19.35257)
INFO     |      - augmented_control_lower_bound(60): 0.00013 (avg 0.00000)
INFO     |      accepted=True ATb_norm=4.82e+00 cost_prev=1258.5291 cost_new=1258.4832
INFO     |  step #80: cost=1190.5166 lambd=0.0640 inexact_tol=2.0e-03
INFO     |      - augmented_control_upper_bound(60): 0.00000 (avg 0.00000)
INFO     |      - augmented_dynamics_constraint(60): 67.96637 (avg 0.28319)
INFO     |      - boundary_cost(12): 29.11151 (avg 0.60649)
INFO     |      - control_cost(60): 1161.40515 (avg 19.35675)
INFO     |      - augmented_control_lower_bound(60): 0.00013 (avg 0.00000)
INFO     |      accepted=True ATb_norm=5.23e+00 cost_prev=1258.4832 cost_new=1258.4417
INFO     |  step #81: cost=1190.6038 lambd=0.0320 inexact_tol=2.0e-03
INFO     |      - augmented_control_upper_bound(60): 0.00000 (avg 0.00000)
INFO     |      - augmented_dynamics_constraint(60): 67.83788 (avg 0.28266)
INFO     |      - boundary_cost(12): 28.95804 (avg 0.60329)
INFO     |      - control_cost(60): 1161.64575 (avg 19.36076)
INFO     |      - augmented_control_lower_bound(60): 0.00013 (avg 0.00000)
INFO     |      accepted=True ATb_norm=4.97e+00 cost_prev=1258.4417 cost_new=1258.4055
INFO     |  step #82: cost=1190.6903 lambd=0.0160 inexact_tol=2.0e-03
INFO     |      - augmented_control_upper_bound(60): 0.00000 (avg 0.00000)
INFO     |      - augmented_dynamics_constraint(60): 67.71500 (avg 0.28215)
INFO     |      - boundary_cost(12): 28.81336 (avg 0.60028)
INFO     |      - control_cost(60): 1161.87695 (avg 19.36462)
INFO     |      - augmented_control_lower_bound(60): 0.00013 (avg 0.00000)
INFO     |      accepted=True ATb_norm=4.56e+00 cost_prev=1258.4055 cost_new=1258.3765
INFO     |  step #83: cost=1190.7764 lambd=0.0080 inexact_tol=2.0e-03
INFO     |      - augmented_control_upper_bound(60): 0.00000 (avg 0.00000)
INFO     |      - augmented_dynamics_constraint(60): 67.59999 (avg 0.28167)
INFO     |      - boundary_cost(12): 28.67562 (avg 0.59741)
INFO     |      - control_cost(60): 1162.10071 (avg 19.36835)
INFO     |      - augmented_control_lower_bound(60): 0.00012 (avg 0.00000)
INFO     |      accepted=True ATb_norm=3.97e+00 cost_prev=1258.3765 cost_new=1258.3521
INFO     |  step #84: cost=1190.8605 lambd=0.0040 inexact_tol=2.0e-03
INFO     |      - augmented_control_upper_bound(60): 0.00000 (avg 0.00000)
INFO     |      - augmented_dynamics_constraint(60): 67.49140 (avg 0.28121)
INFO     |      - boundary_cost(12): 28.54965 (avg 0.59478)
INFO     |      - control_cost(60): 1162.31079 (avg 19.37185)
INFO     |      - augmented_control_lower_bound(60): 0.00012 (avg 0.00000)
INFO     |      accepted=True ATb_norm=3.39e+00 cost_prev=1258.3521 cost_new=1258.3323
INFO     |  step #85: cost=1190.9399 lambd=0.0020 inexact_tol=2.0e-03
INFO     |      - augmented_control_upper_bound(60): 0.00000 (avg 0.00000)
INFO     |      - augmented_dynamics_constraint(60): 67.39226 (avg 0.28080)
INFO     |      - boundary_cost(12): 28.43404 (avg 0.59238)
INFO     |      - control_cost(60): 1162.50586 (avg 19.37510)
INFO     |      - augmented_control_lower_bound(60): 0.00012 (avg 0.00000)
INFO     |      accepted=True ATb_norm=2.91e+00 cost_prev=1258.3323 cost_new=1258.3165
INFO     |  step #86: cost=1191.0142 lambd=0.0010 inexact_tol=2.0e-03
INFO     |      - augmented_control_upper_bound(60): 0.00000 (avg 0.00000)
INFO     |      - augmented_dynamics_constraint(60): 67.30225 (avg 0.28043)
INFO     |      - boundary_cost(12): 28.33015 (avg 0.59021)
INFO     |      - control_cost(60): 1162.68396 (avg 19.37807)
INFO     |      - augmented_control_lower_bound(60): 0.00012 (avg 0.00000)
INFO     |      accepted=True ATb_norm=2.41e+00 cost_prev=1258.3165 cost_new=1258.3042
INFO     |  step #87: cost=1191.0836 lambd=0.0005 inexact_tol=2.0e-03
INFO     |      - augmented_control_upper_bound(60): 0.00000 (avg 0.00000)
INFO     |      - augmented_dynamics_constraint(60): 67.22034 (avg 0.28008)
INFO     |      - boundary_cost(12): 28.23646 (avg 0.58826)
INFO     |      - control_cost(60): 1162.84717 (avg 19.38079)
INFO     |      - augmented_control_lower_bound(60): 0.00012 (avg 0.00000)
INFO     |      accepted=True ATb_norm=2.08e+00 cost_prev=1258.3042 cost_new=1258.2939
INFO     |  step #88: cost=1191.1472 lambd=0.0003 inexact_tol=2.0e-03
INFO     |      - augmented_control_upper_bound(60): 0.00000 (avg 0.00000)
INFO     |      - augmented_dynamics_constraint(60): 67.14668 (avg 0.27978)
INFO     |      - boundary_cost(12): 28.15303 (avg 0.58652)
INFO     |      - control_cost(60): 1162.99414 (avg 19.38324)
INFO     |      - augmented_control_lower_bound(60): 0.00011 (avg 0.00000)
INFO     |      accepted=True ATb_norm=1.77e+00 cost_prev=1258.2939 cost_new=1258.2863
INFO     |  step #89: cost=1191.2051 lambd=0.0001 inexact_tol=2.0e-03
INFO     |      - augmented_control_upper_bound(60): 0.00000 (avg 0.00000)
INFO     |      - augmented_dynamics_constraint(60): 67.08105 (avg 0.27950)
INFO     |      - boundary_cost(12): 28.07865 (avg 0.58497)
INFO     |      - control_cost(60): 1163.12646 (avg 19.38544)
INFO     |      - augmented_control_lower_bound(60): 0.00011 (avg 0.00000)
INFO     |      accepted=True ATb_norm=1.52e+00 cost_prev=1258.2863 cost_new=1258.2808
INFO     |  step #90: cost=1191.2579 lambd=0.0001 inexact_tol=2.0e-03
INFO     |      - augmented_control_upper_bound(60): 0.00000 (avg 0.00000)
INFO     |      - augmented_dynamics_constraint(60): 67.02253 (avg 0.27926)
INFO     |      - boundary_cost(12): 28.01281 (avg 0.58360)
INFO     |      - control_cost(60): 1163.24512 (avg 19.38742)
INFO     |      - augmented_control_lower_bound(60): 0.00011 (avg 0.00000)
INFO     |      accepted=True ATb_norm=1.26e+00 cost_prev=1258.2808 cost_new=1258.2755
INFO     |  step #91: cost=1191.3059 lambd=0.0000 inexact_tol=2.0e-03
INFO     |      - augmented_control_upper_bound(60): 0.00000 (avg 0.00000)
INFO     |      - augmented_dynamics_constraint(60): 66.96944 (avg 0.27904)
INFO     |      - boundary_cost(12): 27.95463 (avg 0.58239)
INFO     |      - control_cost(60): 1163.35132 (avg 19.38919)
INFO     |      - augmented_control_lower_bound(60): 0.00011 (avg 0.00000)
INFO     |      accepted=True ATb_norm=1.05e+00 cost_prev=1258.2755 cost_new=1258.2717
INFO     |  step #92: cost=1191.3483 lambd=0.0000 inexact_tol=2.0e-03
INFO     |      - augmented_control_upper_bound(60): 0.00000 (avg 0.00000)
INFO     |      - augmented_dynamics_constraint(60): 66.92334 (avg 0.27885)
INFO     |      - boundary_cost(12): 27.90324 (avg 0.58132)
INFO     |      - control_cost(60): 1163.44507 (avg 19.39075)
INFO     |      - augmented_control_lower_bound(60): 0.00011 (avg 0.00000)
INFO     |      accepted=True ATb_norm=9.59e-01 cost_prev=1258.2717 cost_new=1258.2689
INFO     |  step #93: cost=1191.3864 lambd=0.0000 inexact_tol=2.0e-03
INFO     |      - augmented_control_upper_bound(60): 0.00000 (avg 0.00000)
INFO     |      - augmented_dynamics_constraint(60): 66.88254 (avg 0.27868)
INFO     |      - boundary_cost(12): 27.85800 (avg 0.58037)
INFO     |      - control_cost(60): 1163.52832 (avg 19.39214)
INFO     |      - augmented_control_lower_bound(60): 0.00011 (avg 0.00000)
INFO     |      accepted=True ATb_norm=7.97e-01 cost_prev=1258.2689 cost_new=1258.2666
INFO     |  step #94: cost=1191.4202 lambd=0.0000 inexact_tol=2.0e-03
INFO     |      - augmented_control_upper_bound(60): 0.00000 (avg 0.00000)
INFO     |      - augmented_dynamics_constraint(60): 66.84635 (avg 0.27853)
INFO     |      - boundary_cost(12): 27.81836 (avg 0.57955)
INFO     |      - control_cost(60): 1163.60181 (avg 19.39336)
INFO     |      - augmented_control_lower_bound(60): 0.00011 (avg 0.00000)
INFO     |      accepted=True ATb_norm=7.15e-01 cost_prev=1258.2666 cost_new=1258.2654
INFO     |  step #95: cost=1191.4502 lambd=0.0000 inexact_tol=2.0e-03
INFO     |      - augmented_control_upper_bound(60): 0.00000 (avg 0.00000)
INFO     |      - augmented_dynamics_constraint(60): 66.81509 (avg 0.27840)
INFO     |      - boundary_cost(12): 27.78339 (avg 0.57882)
INFO     |      - control_cost(60): 1163.66675 (avg 19.39445)
INFO     |      - augmented_control_lower_bound(60): 0.00011 (avg 0.00000)
INFO     |      accepted=True ATb_norm=6.18e-01 cost_prev=1258.2654 cost_new=1258.2643
INFO     |  step #96: cost=1191.4769 lambd=0.0000 inexact_tol=2.0e-03
INFO     |      - augmented_control_upper_bound(60): 0.00000 (avg 0.00000)
INFO     |      - augmented_dynamics_constraint(60): 66.78718 (avg 0.27828)
INFO     |      - boundary_cost(12): 27.75293 (avg 0.57819)
INFO     |      - control_cost(60): 1163.72400 (avg 19.39540)
INFO     |      - augmented_control_lower_bound(60): 0.00011 (avg 0.00000)
INFO     |      accepted=True ATb_norm=5.31e-01 cost_prev=1258.2643 cost_new=1258.2634
INFO     |  step #97: cost=1191.5005 lambd=0.0000 inexact_tol=2.0e-03
INFO     |      - augmented_control_upper_bound(60): 0.00000 (avg 0.00000)
INFO     |      - augmented_dynamics_constraint(60): 66.76284 (avg 0.27818)
INFO     |      - boundary_cost(12): 27.72620 (avg 0.57763)
INFO     |      - control_cost(60): 1163.77429 (avg 19.39624)
INFO     |      - augmented_control_lower_bound(60): 0.00011 (avg 0.00000)
INFO     |      accepted=True ATb_norm=4.97e-01 cost_prev=1258.2634 cost_new=1258.2622
INFO     |  step #98: cost=1191.5215 lambd=0.0000 inexact_tol=2.0e-03
INFO     |      - augmented_control_upper_bound(60): 0.00000 (avg 0.00000)
INFO     |      - augmented_dynamics_constraint(60): 66.74070 (avg 0.27809)
INFO     |      - boundary_cost(12): 27.70294 (avg 0.57714)
INFO     |      - control_cost(60): 1163.81848 (avg 19.39698)
INFO     |      - augmented_control_lower_bound(60): 0.00011 (avg 0.00000)
INFO     |      accepted=True ATb_norm=4.78e-01 cost_prev=1258.2622 cost_new=1258.2617
INFO     |  step #99: cost=1191.5392 lambd=0.0000 inexact_tol=2.0e-03
INFO     |      - augmented_control_upper_bound(60): 0.00000 (avg 0.00000)
INFO     |      - augmented_dynamics_constraint(60): 66.72228 (avg 0.27801)
INFO     |      - boundary_cost(12): 27.68249 (avg 0.57672)
INFO     |      - control_cost(60): 1163.85669 (avg 19.39761)
INFO     |      - augmented_control_lower_bound(60): 0.00011 (avg 0.00000)
INFO     |      accepted=True ATb_norm=4.34e-01 cost_prev=1258.2617 cost_new=1258.2615
INFO     | Terminated @ iteration #100: cost=1191.5549 criteria=[0 0 0], term_deltas=1.3e-05,1.9e-01,4.1e-05 (solved in 0.9474 sec)

Visualization#

Extract solution trajectories and create animated visualization:

# Extract solution trajectories.
states = solution[state_vars]  # Shape: (n_steps+1, 4)
controls = solution[control_vars]  # Shape: (n_steps, 1)

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(states[-1, 0]):.3f}, theta={float(states[-1, 1]):.3f}")
print(f"Target theta (pi): {jnp.pi:.3f}")

Final state: x=0.000, theta=3.141
Target theta (pi): 3.142

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."""
    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(states[i, 0]), float(states[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(states[0, 0]), float(states[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 (Collocation)",
        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 states[:, 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 states[:, 1]],
        mode="lines",
        line=dict(color="coral", width=2),
        name="θ",
    ),
    row=1,
    col=2,
)
fig_traj.add_hline(
    y=float(jnp.pi),
    line_dash="dash",
    line_color="gray",
    row=1,
    col=2,
    annotation_text="π",
)

# Velocities.
fig_traj.add_trace(
    go.Scatter(
        x=[float(t) for t in times],
        y=[float(s) for s in states[:, 2]],
        mode="lines",
        line=dict(color="steelblue", width=2),
        name="ẋ",
    ),
    row=2,
    col=1,
)
fig_traj.add_trace(
    go.Scatter(
        x=[float(t) for t in times],
        y=[float(s) for s in states[:, 3]],
        mode="lines",
        line=dict(color="coral", width=2, dash="dash"),
        name="θ̇",
    ),
    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 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,
)
fig_traj.add_hline(
    y=force_limit, line_dash="dash", line_color="red", opacity=0.5, row=2, col=2
)
fig_traj.add_hline(
    y=-force_limit, line_dash="dash", line_color="red", opacity=0.5, 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"))
Contents