Qube nn

SPDX-License-Identifier: MIT

In [ ]:

import pickle

import pickle

In [ ]:

import matplotlib.pyplot as plt

import matplotlib.pyplot as plt

In [ ]:

import numpy as np
import jax
import jax.numpy as jnp

import numpy as np import jax import jax.numpy as jnp

In [ ]:

import jaxonomy
from jaxonomy import library
from jaxonomy.optimization import Trainer

import jaxonomy from jaxonomy import library from jaxonomy.optimization import Trainer

In [ ]:

from controllers import make_mlp_controller

from controllers import make_mlp_controller

In [ ]:

RTOL = 1e-6
ATOL = 1e-8

RTOL = 1e-6 ATOL = 1e-8

In [ ]:

def make_cost_system(Q, R, name="cost"):
    builder = jaxonomy.DiagramBuilder()
    lqr_cost = library.QuadraticCost(Q, R, name="lqr_cost")
    running_cost = library.Integrator(initial_state=0.0, name="running_cost")
    builder.add(lqr_cost, running_cost)
    builder.connect(lqr_cost.output_ports[0], running_cost.input_ports[0])
    builder.export_input(lqr_cost.input_ports[0], name="x")
    builder.export_input(lqr_cost.input_ports[1], name="u")
    builder.export_output(running_cost.output_ports[0])
    return builder.build(name=name)

def make_cost_system(Q, R, name="cost"): builder = jaxonomy.DiagramBuilder() lqr_cost = library.QuadraticCost(Q, R, name="lqr_cost") running_cost = library.Integrator(initial_state=0.0, name="running_cost") builder.add(lqr_cost, running_cost) builder.connect(lqr_cost.output_ports[0], running_cost.input_ports[0]) builder.export_input(lqr_cost.input_ports[0], name="x") builder.export_input(lqr_cost.input_ports[1], name="u") builder.export_output(running_cost.output_ports[0]) return builder.build(name=name)

In [ ]:

# Full closed-loop system
def make_cl_system(
    plant,
    nn_config,
    reference=None,
    dt=1.0,
    Q=1.0,
    R=1e-2,
    lbu=-np.inf,
    ubu=np.inf,
    sigma=0.0,
    filter_coefficient=100.0,
    delay=0.0,
    name="cl_system",
):
    builder = jaxonomy.DiagramBuilder()

    plant = builder.add(plant)

    # Add a reference input which will also be fed to the NN controller
    if reference is None:
        reference = 0.0

    reference = library.Constant(reference, name="reference")
    builder.add(reference)

    error = builder.add(library.Adder(2, operators="-+", name="error"))
    builder.connect(plant.output_ports[0], error.input_ports[0])
    builder.connect(reference.output_ports[0], error.input_ports[1])

    # Create the NN controller
    controller = builder.add(
        make_mlp_controller(
            nn_config,
            dt=dt,
            sigma=sigma,
            filter_coefficient=filter_coefficient,
            delay=delay,
            name="controller",
        )
    )

    # Combine the reference and plant output for the controller inputs
    builder.connect(error.output_ports[0], controller.input_ports[1])
    saturate = builder.add(
        library.Saturate(lower_limit=lbu, upper_limit=ubu, name="saturate")
    )
    builder.connect(controller.output_ports[0], saturate.input_ports[0])
    builder.connect(saturate.output_ports[0], plant.input_ports[0])

    # Feed the saturated control signal back to the controller
    # This is unused in this case, but is needed for compatibility with
    # other controllers that use, e.g. EKF estimators
    builder.connect(saturate.output_ports[0], controller.input_ports[0])

    # Cost-tracking subsystem
    cost = builder.add(make_cost_system(Q, R, name="cost"))
    builder.connect(controller.output_ports[1], cost.input_ports[0])
    builder.connect(saturate.output_ports[0], cost.input_ports[1])

    return builder.build(name=name)

# Full closed-loop system def make_cl_system( plant, nn_config, reference=None, dt=1.0, Q=1.0, R=1e-2, lbu=-np.inf, ubu=np.inf, sigma=0.0, filter_coefficient=100.0, delay=0.0, name="cl_system", ): builder = jaxonomy.DiagramBuilder() plant = builder.add(plant) # Add a reference input which will also be fed to the NN controller if reference is None: reference = 0.0 reference = library.Constant(reference, name="reference") builder.add(reference) error = builder.add(library.Adder(2, operators="-+", name="error")) builder.connect(plant.output_ports[0], error.input_ports[0]) builder.connect(reference.output_ports[0], error.input_ports[1]) # Create the NN controller controller = builder.add( make_mlp_controller( nn_config, dt=dt, sigma=sigma, filter_coefficient=filter_coefficient, delay=delay, name="controller", ) ) # Combine the reference and plant output for the controller inputs builder.connect(error.output_ports[0], controller.input_ports[1]) saturate = builder.add( library.Saturate(lower_limit=lbu, upper_limit=ubu, name="saturate") ) builder.connect(controller.output_ports[0], saturate.input_ports[0]) builder.connect(saturate.output_ports[0], plant.input_ports[0]) # Feed the saturated control signal back to the controller # This is unused in this case, but is needed for compatibility with # other controllers that use, e.g. EKF estimators builder.connect(saturate.output_ports[0], controller.input_ports[0]) # Cost-tracking subsystem cost = builder.add(make_cost_system(Q, R, name="cost")) builder.connect(controller.output_ports[1], cost.input_ports[0]) builder.connect(saturate.output_ports[0], cost.input_ports[1]) return builder.build(name=name)

In [ ]:

class QubeTrainer(Trainer):
    def __init__(self, *args, nn_id, ref_id, plant_id, **kwargs):
        self.nn_id = nn_id
        self.ref_id = ref_id
        self.plant_id = plant_id
        super().__init__(*args, **kwargs)

    # Model-specific: what parameters can we optimize?
    def optimizable_parameters(self, context):
        return context[self.nn_id].parameters["mlp_params"]

    # Model-specific updates to add the sample data, parameters, etc.
    # This should be the combination of the output of `optimizable_parameters`
    # along with all the per-simulation "training data".  Parameters will
    # update once per epoch, and training data will update once per sample.
    def prepare_context(self, context, mlp_params, x0, key):
        nn_context = context[self.nn_id].with_parameter("mlp_params", mlp_params)
        plant_context = context[self.plant_id].with_continuous_state(x0)
        noise = self.simulator.system["controller"]["noise"]

        noise_context = context[noise.system_id].with_discrete_state(
            noise.RNGState(key, noise.sample(key, shape=(noise.N, *noise.shape)))
        )
        context = (
            context.with_subcontext(self.nn_id, nn_context)
            .with_subcontext(self.plant_id, plant_context)
            .with_subcontext(noise.system_id, noise_context)
        )
        return context

    # Model-specific cost function
    def evaluate_cost(self, context):
        system = self.simulator.system
        T = context.time
        J = (1 / T) * system["cost"].output_ports[0].eval(context)
        return J

class QubeTrainer(Trainer): def __init__(self, *args, nn_id, ref_id, plant_id, **kwargs): self.nn_id = nn_id self.ref_id = ref_id self.plant_id = plant_id super().__init__(*args, **kwargs) # Model-specific: what parameters can we optimize? def optimizable_parameters(self, context): return context[self.nn_id].parameters["mlp_params"] # Model-specific updates to add the sample data, parameters, etc. # This should be the combination of the output of `optimizable_parameters` # along with all the per-simulation "training data". Parameters will # update once per epoch, and training data will update once per sample. def prepare_context(self, context, mlp_params, x0, key): nn_context = context[self.nn_id].with_parameter("mlp_params", mlp_params) plant_context = context[self.plant_id].with_continuous_state(x0) noise = self.simulator.system["controller"]["noise"] noise_context = context[noise.system_id].with_discrete_state( noise.RNGState(key, noise.sample(key, shape=(noise.N, *noise.shape))) ) context = ( context.with_subcontext(self.nn_id, nn_context) .with_subcontext(self.plant_id, plant_context) .with_subcontext(noise.system_id, noise_context) ) return context # Model-specific cost function def evaluate_cost(self, context): system = self.simulator.system T = context.time J = (1 / T) * system["cost"].output_ports[0].eval(context) return J

In [ ]:

def plot_rollout(results):
    fig, ax = plt.subplots(3, 1, figsize=(7, 3), sharex=True)
    ax[-1].set_xlabel("$t$")

    ax[0].plot(results.time, results.outputs["x"][:, 1])
    ax[0].plot(results.time, results.outputs["r"][:, 1], "k")
    ax[0].set_ylabel("$x$")
    ax[0].grid()

    ax[1].plot(results.time, results.outputs["u"])
    ax[1].set_ylabel("$u$")
    ax[1].grid()

    ax[2].plot(results.time, results.outputs["J"])
    ax[2].set_ylabel("$J$")
    ax[2].grid()
    plt.show()

def plot_rollout(results): fig, ax = plt.subplots(3, 1, figsize=(7, 3), sharex=True) ax[-1].set_xlabel("$t$") ax[0].plot(results.time, results.outputs["x"][:, 1]) ax[0].plot(results.time, results.outputs["r"][:, 1], "k") ax[0].set_ylabel("$x$") ax[0].grid() ax[1].plot(results.time, results.outputs["u"]) ax[1].set_ylabel("$u$") ax[1].grid() ax[2].plot(results.time, results.outputs["J"]) ax[2].set_ylabel("$J$") ax[2].grid() plt.show()

In [ ]:

def run_rollout(nn_config, sys_config, p_opt, plant=None, x0=None, tf=10.0, plot=True):
    if plant is None:
        plant = library.QubeServoModel(full_state_output=False, name="plant")

    # Simulate the system with the trained controller
    system = make_cl_system(plant, nn_config, **sys_config)
    context = system.create_context()

    nn_id = system["controller"]["net"].system_id
    nn_context = context[nn_id].with_parameter("mlp_params", p_opt)
    context = context.with_subcontext(nn_id, nn_context)

    if x0 is not None:
        plant_id = system["plant"].system_id
        plant_context = context[plant_id].with_continuous_state(x0)
        context = context.with_subcontext(plant_id, plant_context)

    options = jaxonomy.SimulatorOptions(
        max_major_steps=1000,
        save_time_series=True,
        max_major_step_length=1.0,
        rtol=1e-6,
        atol=1e-8,
    )
    recorded_signals = {
        "x": system["plant"].output_ports[0],
        "r": system["reference"].output_ports[0],
        "u": system["plant"].input_ports[0],
        "u0": system["controller"]["net"].output_ports[0],
        "J": system["cost"].output_ports[0],
        "e": system["error"].output_ports[0],
        "w": system["controller"]["noisy_ctrl"].output_ports[0],
    }

    results = jaxonomy.simulate(
        system, context, (0.0, tf), options=options, recorded_signals=recorded_signals
    )

    if plot:
        plot_rollout(results)

    return results

def run_rollout(nn_config, sys_config, p_opt, plant=None, x0=None, tf=10.0, plot=True): if plant is None: plant = library.QubeServoModel(full_state_output=False, name="plant") # Simulate the system with the trained controller system = make_cl_system(plant, nn_config, **sys_config) context = system.create_context() nn_id = system["controller"]["net"].system_id nn_context = context[nn_id].with_parameter("mlp_params", p_opt) context = context.with_subcontext(nn_id, nn_context) if x0 is not None: plant_id = system["plant"].system_id plant_context = context[plant_id].with_continuous_state(x0) context = context.with_subcontext(plant_id, plant_context) options = jaxonomy.SimulatorOptions( max_major_steps=1000, save_time_series=True, max_major_step_length=1.0, rtol=1e-6, atol=1e-8, ) recorded_signals = { "x": system["plant"].output_ports[0], "r": system["reference"].output_ports[0], "u": system["plant"].input_ports[0], "u0": system["controller"]["net"].output_ports[0], "J": system["cost"].output_ports[0], "e": system["error"].output_ports[0], "w": system["controller"]["noisy_ctrl"].output_ports[0], } results = jaxonomy.simulate( system, context, (0.0, tf), options=options, recorded_signals=recorded_signals ) if plot: plot_rollout(results) return results

In [ ]:

# Run a full train-rollout experiment
def run_experiment(
    nn_config,
    sys_config,
    tf=10.0,
    lr=1e-2,
    N=1000,
    epochs=50,
    batch_size=None,
    x0_std=1.0,
    x0=None,
    key=None,
    init_params=None,
    opt_state=None,
):
    plant = library.QubeServoModel(full_state_output=False, name="plant")
    system = make_cl_system(plant, nn_config, **sys_config)
    context = system.create_context()

    # Generate random initial conditions
    if key is None:
        key = jax.random.PRNGKey(seed=nn_config["seed"])

    if x0 is None:
        key, subkey = jax.random.split(key)
        x_up = np.array([0.0, np.pi, 0.0, 0.0])  # Up position as reference
        x0 = x_up + x0_std * jax.random.normal(subkey, (N, 4))

    if batch_size is None:
        batch_size = N

    # Reshape the initial conditions into a batch
    x0 = jnp.reshape(x0, (-1, batch_size, *x0.shape[1:]))

    training_data = (x0,)

    options = jaxonomy.SimulatorOptions(
        max_major_steps=200,
        enable_autodiff=True,
        rtol=1e-6,
        atol=1e-8,
    )
    sim = jaxonomy.Simulator(system, options=options)
    trainer = QubeTrainer(
        sim,
        context,
        nn_id=system["controller"]["net"].system_id,
        ref_id=system["reference"].system_id,
        plant_id=system["plant"].system_id,
        lr=lr,
        optimizer="adam",
    )
    opt_params = trainer.train(
        training_data,
        sim_start_time=0.0,
        sim_stop_time=tf,
        epochs=epochs,
        key=key,
        params=init_params,
        opt_state=opt_state,
    )

    return opt_params, trainer.opt_state

# Run a full train-rollout experiment def run_experiment( nn_config, sys_config, tf=10.0, lr=1e-2, N=1000, epochs=50, batch_size=None, x0_std=1.0, x0=None, key=None, init_params=None, opt_state=None, ): plant = library.QubeServoModel(full_state_output=False, name="plant") system = make_cl_system(plant, nn_config, **sys_config) context = system.create_context() # Generate random initial conditions if key is None: key = jax.random.PRNGKey(seed=nn_config["seed"]) if x0 is None: key, subkey = jax.random.split(key) x_up = np.array([0.0, np.pi, 0.0, 0.0]) # Up position as reference x0 = x_up + x0_std * jax.random.normal(subkey, (N, 4)) if batch_size is None: batch_size = N # Reshape the initial conditions into a batch x0 = jnp.reshape(x0, (-1, batch_size, *x0.shape[1:])) training_data = (x0,) options = jaxonomy.SimulatorOptions( max_major_steps=200, enable_autodiff=True, rtol=1e-6, atol=1e-8, ) sim = jaxonomy.Simulator(system, options=options) trainer = QubeTrainer( sim, context, nn_id=system["controller"]["net"].system_id, ref_id=system["reference"].system_id, plant_id=system["plant"].system_id, lr=lr, optimizer="adam", ) opt_params = trainer.train( training_data, sim_start_time=0.0, sim_stop_time=tf, epochs=epochs, key=key, params=init_params, opt_state=opt_state, ) return opt_params, trainer.opt_state

In [ ]:

if __name__ == "__main__":
    nn_config = {
        "seed": 0,
        "in_size": 4,
        "out_size": 1,
        "width_size": 16,
        "depth": 4,
        "activation_str": "swish",
    }

    # Weight the states higher than the controls so that the network doesn't
    # learn a suboptimal strategy of just leaving the pendulum in the down
    # position
    Q = np.diag([1.0, 1.0, 1e-2, 1e-2])
    R = 1e-2 * np.eye(1)
    dt = 1 / 500
    sys_config = {
        "dt": dt,
        "Q": Q,
        "R": R,
        "lbu": -10,
        "ubu": 10,
        "reference": np.array([0.0, np.pi]),
        "sigma": 1e-2,
    }

    tf = 0.5  # Time horizon for training
    N = 100  # Number of random initial conditions
    x0 = np.zeros((N, 4)) + np.random.randn(N, 4) * 0.1
    key = jax.random.PRNGKey(0)
    opt_params, opt_state = run_experiment(
        nn_config=nn_config,
        sys_config=sys_config,
        tf=tf,
        lr=1e-3,
        epochs=1000,
        batch_size=10,
        x0=x0,
        key=key,
    )

    with open("models/swingup.pkl", "wb") as f:
        pickle.dump((opt_params, opt_state, nn_config, sys_config), f)

if __name__ == "__main__": nn_config = { "seed": 0, "in_size": 4, "out_size": 1, "width_size": 16, "depth": 4, "activation_str": "swish", } # Weight the states higher than the controls so that the network doesn't # learn a suboptimal strategy of just leaving the pendulum in the down # position Q = np.diag([1.0, 1.0, 1e-2, 1e-2]) R = 1e-2 * np.eye(1) dt = 1 / 500 sys_config = { "dt": dt, "Q": Q, "R": R, "lbu": -10, "ubu": 10, "reference": np.array([0.0, np.pi]), "sigma": 1e-2, } tf = 0.5 # Time horizon for training N = 100 # Number of random initial conditions x0 = np.zeros((N, 4)) + np.random.randn(N, 4) * 0.1 key = jax.random.PRNGKey(0) opt_params, opt_state = run_experiment( nn_config=nn_config, sys_config=sys_config, tf=tf, lr=1e-3, epochs=1000, batch_size=10, x0=x0, key=key, ) with open("models/swingup.pkl", "wb") as f: pickle.dump((opt_params, opt_state, nn_config, sys_config), f)