04 trajopt

In [ ]:

import os
import sys
from time import sleep

import os import sys from time import sleep

In [ ]:

import numpy as np

import numpy as np

In [ ]:

import jax
import jax.numpy as jnp

import jax import jax.numpy as jnp

In [ ]:

import collimator
from collimator import library

import collimator from collimator import library

In [ ]:

from controllers import *  # noqa: F403

from controllers import * # noqa: F403

In [ ]:

HOME = os.path.expanduser("~")
sys.path.append(f"{HOME}/Quanser/libraries/python")

HOME = os.path.expanduser("~") sys.path.append(f"{HOME}/Quanser/libraries/python")

In [ ]:

#
# Design the optimal trajectory
#
plant = library.QubeServoModel(full_state_output=True, name="qube")

# # Design the optimal trajectory # plant = library.QubeServoModel(full_state_output=True, name="qube")

In [ ]:

nx, nu, ny = 4, 1, 2  # State, control, measurement dimensions
t0, th = 0.0, 1.0  # Initial and final times
x0 = np.array([0.0, 0.0, 0.0, 0.0])  # Initial state
xf = np.array([0.0, np.pi, 0.0, 0.0])  # Target state
N = 50  # Discretization points

nx, nu, ny = 4, 1, 2 # State, control, measurement dimensions t0, th = 0.0, 1.0 # Initial and final times x0 = np.array([0.0, 0.0, 0.0, 0.0]) # Initial state xf = np.array([0.0, np.pi, 0.0, 0.0]) # Target state N = 50 # Discretization points

In [ ]:

# Optimize the swing-up trajectory, constraining only the final state
t_opt = np.linspace(t0, th, N + 1)
print("Solving for optimal trajectory...")
x_opt, u_opt = collimator.trajopt(
    plant,
    t0=t0,
    tf=th,
    x0=x0,
    xf=xf,
    Q=1e-2 * np.eye(nx),
    R=np.eye(nu),
    N=N,
    constrain_xf=True,
    lb_u=-15.0,
    ub_u=15.0,
)

# Optimize the swing-up trajectory, constraining only the final state t_opt = np.linspace(t0, th, N + 1) print("Solving for optimal trajectory...") x_opt, u_opt = collimator.trajopt( plant, t0=t0, tf=th, x0=x0, xf=xf, Q=1e-2 * np.eye(nx), R=np.eye(nu), N=N, constrain_xf=True, lb_u=-15.0, ub_u=15.0, )

In [ ]:

interp_fun = jax.vmap(jnp.interp, (None, None, 1))

interp_fun = jax.vmap(jnp.interp, (None, None, 1))

In [ ]:

# Offset by at least 2 seconds to make sure the timing is aligned
t_offset = 2.0

# Offset by at least 2 seconds to make sure the timing is aligned t_offset = 2.0

In [ ]:

def nominal_trajectory_x(t):
    tp = t - t_offset
    return jnp.where(
        (tp > th), xf, jnp.where(tp < 0.0, 0.0, interp_fun(tp, t_opt, x_opt))
    )

def nominal_trajectory_x(t): tp = t - t_offset return jnp.where( (tp > th), xf, jnp.where(tp < 0.0, 0.0, interp_fun(tp, t_opt, x_opt)) )

In [ ]:

def nominal_trajectory_u(t):
    tp = t - t_offset
    return jnp.where((tp > th) | (tp < 0.0), 0.0, interp_fun(tp, t_opt, u_opt))

def nominal_trajectory_u(t): tp = t - t_offset return jnp.where((tp > th) | (tp < 0.0), 0.0, interp_fun(tp, t_opt, u_opt))

In [ ]:

def make_ffwd_swingup(dt, t0, tf, noise_amplitude=1e2, R=1e0, name="ffwd"):
    # LQR parameters
    Q = np.eye(nx)  # state penalty matrix
    # Q = np.diag([0.0, 1.0, 0.0, 1.0])
    Qf = Q  # terminal state penalty matrix
    N = np.zeros((nx, nu))  # cross cost matrix between state and control vectors

    # Kalman filter parameters
    G = np.eye(nu)
    QN = noise_amplitude * np.eye(nu)
    RN = np.eye(ny)

    ekf = library.ExtendedKalmanFilter.for_continuous_plant(
        library.QubeServoModel(full_state_output=False),
        dt=dt,
        G_func=lambda t: G,
        Q_func=lambda t: QN,
        R_func=lambda t: RN,
        x_hat_0=nominal_trajectory_x(0),
        P_hat_0=1e-4 * np.eye(nx),
        discretized_noise=True,
        name="ekf",
    )

    # create finite-horizon LQR controller with identical nominal and desired trajectories
    lqr = library.FiniteHorizonLinearQuadraticRegulator(
        t0,
        tf,
        library.QubeServoModel(full_state_output=True, name="qube"),
        Qf,
        func_Q=lambda t: Q,
        func_R=lambda t: R * np.eye(nu),
        func_N=lambda t: N,
        func_x_0=nominal_trajectory_x,
        func_u_0=nominal_trajectory_u,
        name="lqr",
    )

    builder = collimator.DiagramBuilder()
    builder.add(ekf, lqr)
    builder.connect(ekf.output_ports[0], lqr.input_ports[0])

    # Export inputs for (u, y)
    builder.export_input(ekf.input_ports[0], "u")
    builder.export_input(ekf.input_ports[1], "y")

    # Export outputs for (u, x_hat)
    builder.export_output(lqr.output_ports[0], "u")
    builder.export_output(ekf.output_ports[0], "x_hat")

    return builder.build(name=name)

def make_ffwd_swingup(dt, t0, tf, noise_amplitude=1e2, R=1e0, name="ffwd"): # LQR parameters Q = np.eye(nx) # state penalty matrix # Q = np.diag([0.0, 1.0, 0.0, 1.0]) Qf = Q # terminal state penalty matrix N = np.zeros((nx, nu)) # cross cost matrix between state and control vectors # Kalman filter parameters G = np.eye(nu) QN = noise_amplitude * np.eye(nu) RN = np.eye(ny) ekf = library.ExtendedKalmanFilter.for_continuous_plant( library.QubeServoModel(full_state_output=False), dt=dt, G_func=lambda t: G, Q_func=lambda t: QN, R_func=lambda t: RN, x_hat_0=nominal_trajectory_x(0), P_hat_0=1e-4 * np.eye(nx), discretized_noise=True, name="ekf", ) # create finite-horizon LQR controller with identical nominal and desired trajectories lqr = library.FiniteHorizonLinearQuadraticRegulator( t0, tf, library.QubeServoModel(full_state_output=True, name="qube"), Qf, func_Q=lambda t: Q, func_R=lambda t: R * np.eye(nu), func_N=lambda t: N, func_x_0=nominal_trajectory_x, func_u_0=nominal_trajectory_u, name="lqr", ) builder = collimator.DiagramBuilder() builder.add(ekf, lqr) builder.connect(ekf.output_ports[0], lqr.input_ports[0]) # Export inputs for (u, y) builder.export_input(ekf.input_ports[0], "u") builder.export_input(ekf.input_ports[1], "y") # Export outputs for (u, x_hat) builder.export_output(lqr.output_ports[0], "u") builder.export_output(ekf.output_ports[0], "x_hat") return builder.build(name=name)

In [ ]:

def make_diagram(plant, dt, t0, tf, noise_amplitude=1e2, R=1e0):
    swingup_controller = make_ffwd_swingup(dt, t0, tf, noise_amplitude, R)
    balance_controller = make_pid(dt, y_eq=np.array([0.0, np.pi]))  # noqa: F405

    controller = make_switched_controller(  # noqa: F405
        swingup_controller, balance_controller, threshold=0.25
    )

    builder = collimator.DiagramBuilder()
    builder.add(plant, controller)
    builder.connect(plant.output_ports[0], controller.input_ports[0])
    builder.connect(controller.output_ports[0], plant.input_ports[0])

    return builder.build()

def make_diagram(plant, dt, t0, tf, noise_amplitude=1e2, R=1e0): swingup_controller = make_ffwd_swingup(dt, t0, tf, noise_amplitude, R) balance_controller = make_pid(dt, y_eq=np.array([0.0, np.pi])) # noqa: F405 controller = make_switched_controller( # noqa: F405 swingup_controller, balance_controller, threshold=0.25 ) builder = collimator.DiagramBuilder() builder.add(plant, controller) builder.connect(plant.output_ports[0], controller.input_ports[0]) builder.connect(controller.output_ports[0], plant.input_ports[0]) return builder.build()

In [ ]:

#
# Test the feedforward controller
#
dt = 0.001  # 1 kHz control loop
t0 = 0.0
th = t_opt[-1] + t_offset  # Feedforward time horizon

# # Test the feedforward controller # dt = 0.001 # 1 kHz control loop t0 = 0.0 th = t_opt[-1] + t_offset # Feedforward time horizon

In [ ]:

plant_hw = library.QuanserHAL(dt=dt, version=3, hardware=True)
system = make_diagram(plant_hw, dt, t0, th, R=1e1)
context = system.create_context()

plant_hw = library.QuanserHAL(dt=dt, version=3, hardware=True) system = make_diagram(plant_hw, dt, t0, th, R=1e1) context = system.create_context()

In [ ]:

# Wait for one second to make sure the communication channel is open
sleep(1.0)

# Wait for one second to make sure the communication channel is open sleep(1.0)

In [ ]:

tf = 6.0
results = collimator.simulate(system, context, (0.0, tf))
plant_hw.terminate()

tf = 6.0 results = collimator.simulate(system, context, (0.0, tf)) plant_hw.terminate()