RobustNeuralNetworks.jl: a Package for Machine Learning and Data-Driven Control with Certified Robustness

Firstly, all of our RENs are contracting systems. This means that they exponentially “forget” their initial conditions. If the system starts at two different initial conditions but is given the same input sequence, the internal states will exponentially converge over time. Figure 2 shows an example of a contracting REN with one input and a single internal state, where two simulations of the system start with different initial conditions but are provided the same sinusoidal input. See [Bullo2022] for a detailed introduction to contraction theory for dynamical systems.

We define additional robustness criteria on the input-output map of RENs with incremental integral quadratic constraints (IQCs). Suppose we have a model $\mathcal{M}$ starting at two different initial conditions $a,b$ with two different input signals $u,v$, and consider their corresponding output trajectories $y^{a}=\mathcal{M}_{a}(u)$ and $y^{b}=\mathcal{M}_{b}(v).$ The model $\mathcal{M}$ satisfies the IQC defined by matrices $(Q,S,R)$ if

for some function $d(a,b)\geq 0$ with $d(a,a)=0$, where $0\preceq Q\in\mathbb{R}^{n_{y}\times n_{y}}$, $S\in\mathbb{R}^{n_{u}\times n_{y}},$ $R=R^{\top}\in\mathbb{R}^{n_{u}\times n_{u}}.$

In general, the IQC matrices $(Q,S,R)$ can be chosen (or optimized) to meet a range of performance criteria. There are a few special cases that are worth noting.

If $Q=-\frac{1}{\gamma}I$, $R=\gamma I$, $S=0$ for some $\gamma\in\mathbb{R}$ with $\gamma\neq 0$, the model $\mathcal{M}$ satisfies a Lipschitz bound (incremental $\ell^{2}$-gain bound) of $\gamma$, defined by

Qualitatively, the Lipschitz bound is a measure of how smooth the network is. If the Lipschitz bound $\gamma$ is small, then small changes in the inputs $u,v$ will lead to small changes in the model output. If $\gamma$ is large, then the model output might change significantly for even small changes to the inputs. This can make the model more sensitive to noise, adversarial attacks, and other input disturbances.

As the name suggests, all LBDN models are all constructed to have a user-tunable (or learnable) Lipschitz bound.

We have implemented two versions of incremental passivity. In each case, the network must have the same number of inputs and outputs.

1. 1.

   If $Q=0,R=-2\nu I,S=I$ where $\nu\geq 0$, the model is incrementally passive (incrementally strictly input passive if $\nu>0$). Mathematically, the following inequality holds.

   $$\langle\mathcal{M}_{a}(u)-\mathcal{M}_{b}(v),u-v\rangle_{T}\geq\nu\|u-v\|^{2}_{T}$$

   (11)

2. 2.

   If $Q=-2\rho I,R=0,S=I$ where $\rho>0$, the model is incrementally strictly output passive. Mathematically, the following inequality holds.

   $$\langle\mathcal{M}_{a}(u)-\mathcal{M}_{b}(v),u-v\rangle_{T}\geq\rho\|\mathcal{M}_{a}(u)-\mathcal{M}_{b}(v)\|^{2}_{T}$$

   (12)

For more details on IQCs and their use in RENs, please see [Revay++2021b].

RENs are defined by two abstract types in RobustNeuralNetworks.jl. Subtypes of AbstractRENParams hold all the information required to directly parameterize a REN satisfying some robustness properties. For example, to initialise the direct parameters of a contracting REN with 1 input, 10 states, 20 neurons, 1 output, and a relu activation function, we use the following. The direct parameters $\theta$ are stored in model_ps.direct.

[language = Julia] using Flux, RobustNeuralNetworks

T = Float32 nu, nx, nv, ny = 1, 10, 20, 1 model_ps = ContractingRENParamsT( nu, nx, nv, ny; nl=Flux.relu)

println(model_ps.direct) # Access direct params

Subtypes of AbstractREN represent RENs in their explicit form which can be evaluated on data. The conversion from direct to explicit parameters $\theta\mapsto\bar{\theta}$ is performed when the REN is constructed and the explicit parameters $\bar{\theta}$ are stored in model.explicit.

[language = Julia] model = REN(ps) # Create explicit model println(model.explicit) # Access explicit params

Figure 3 illustrates this architecture. We use a similar interface based on AbstractLBDNParams and AbstractLBDN for LBDNs.

There are currently four REN parameterizations implemented in this package:

1. 1.

   ContractingRENParams parameterizes contracting RENs with a user-defined upper bound on the contraction rate.

2. 2.

   LipschitzRENParams parameterizes RENs with a user-defined (or learnable) Lipschitz bound of $\gamma\in(0,\infty)$.

3. 3.

   PassiveRENParams parameterizes incrementally input passive RENs with user-tunable passivity parameter $\nu\geq 0$.

4. 4.

   GeneralRENParams parameterizes RENs satisfying some general behavioural constraints defined by an incremental IQC with parameters (Q,S,R).

There is currently only one LBDN parameterization implemented in RobustNeuralNetworks.jl:

1. 1.

   DenseLBDNParams parameterizes dense (fully-connected) LBDNs with a user-defined or learnable Lipschitz bound. A dense LBDN is effectively a Lipschitz-bounded MLP.

We intend to add ConvolutionalLBDNParams to parameterize the convolutional LBDNs in [Wang+Manchester2023] in future iterations of the package.

When training a REN or LBDN, we learn and update the direct parameters $\theta$ and convert them to the explicit parameters $\bar{\theta}$ only for model evaluation. The main constructors for explicit models are REN and LBDN.

Users familiar with Flux.jl will be used to creating a model once and then training it on their data. The typical workflow is as follows.

[language = Julia] using Flux

# Define a model and a loss function model = Flux.Chain( Flux.Dense(1 => 10, Flux.relu), Flux.Dense(10 => 1, Flux.relu) )

loss(model, x, y) = Flux.mse(model(x), y)

# Training data of 20 batches T = Float32 xs, ys = rand(T,1,20), rand(T,1,20) data = [(xs, ys)]

# Train the model for 50 epochs opt_state = Flux.setup(Adam(0.01), model) for _in 1:50 Flux.train!(loss, model, data, opt_state) end

When training a model constructed from REN or LBDN, we need to back-propagate through the mapping from direct (learnable) parameters to the explicit model. We must therefore include the model construction as part of the loss function. If we do not, then the auto-differentiation engine has no knowledge of how the model parameters affect the loss, and will return zero gradients. Here is an example with an LBDN, where the model is defined by the direct parameterization stored in model_ps.

[language = Julia] using Flux, RobustNeuralNetworks

# Define model parameterization and loss function T = Float32 model_ps = DenseLBDNParamsT(1, [10], 1; nl=relu) function loss(model_ps, x, y) model = LBDN(model_ps) Flux.mse(model(x), y) end

# Training data of 20 batches xs, ys = rand(T,1,20), rand(T,1,20) data = [(xs, ys)]

# Train the model for 50 epochs opt_state = Flux.setup(Adam(0.01), model_ps) for _in 1:50 Flux.train!(loss, model_ps, data, opt_state) end

For the sake of convenience, we have included the model wrappers DiffREN and DiffLBDN as alternatives to REN and LBDN, respectively. These wrappers compute the explicit parameters each time the model is called rather than just once when they are constructed. Any model created with these wrappers can therefore be used exactly the same way as a regular Flux.jl model, and there is no need for model construction in the loss function. One can simply replace the definition of the Flux.Chain model in the first code cell above with {lstlisting}[language = Julia] model_ps = DenseLBDNParamsT(1, [10], 1; nl=relu) model = DiffLBDN(model_ps) and train the LBDN just like any other Flux.jl model. We use these wrappers in Sections 3.1 and LABEL:sec:observer.

We must, however, caution against their use in problems where a model is evaluated many times before being updated. The model_ps and model are kept separate with REN and LBDN to offer flexibility. The main computational bottleneck in training a REN or LBDN is converting from the direct to explicit parameters (mapping $\bar{f}:\theta\mapsto\bar{\theta}$). For both model structures, this process involves a matrix inverse, where the number of matrix elements scales quadratically with the dimension of the model in a REN and the dimension of each layer in an LBDN (see [Revay++2021b, Wang+Manchester2023]).

As per Section 2.3.1, direct parameters are stored in model_ps, while explicit parameters are computed when the model is created and are stored within it. In some applications (eg: reinforcement learning), a model is called many times with the same explicit parameters $\bar{\theta}$ before its learnable parameters $\theta$ are updated. It is therefore significantly more efficient to store the explicit parameters, use them many times, and then update them only when the learnable parameters change. We cannot store the direct and explicit parameters in the same model object since auto-differentiation in Julia does not support array mutation [Innes2018b]. Instead, we separate the two.

The computational benefits of separating models from their parameterizations is explored numerically in Section 3.2.

We begin by loading the training and test data. MLDatasets.jl²²2https://juliaml.github.io/MLDatasets.jl/ contains a number of common machine-learning datasets, including the MNIST dataset. To load the full dataset of 60,000 training images and 10,000 test images, one would run the following code.

[language = Julia] using MLDatasets: MNIST

T = Float32 x_train, y_train = MNIST(T, split=:train)[:] x_test, y_test = MNIST(T, split=:test)[:]

The feature matrices x_train and x_test are three-dimensional arrays where each 28x28 layer contains pixel data for a single handwritten number from 0 to 9 (eg: see Fig. 4). The labels y_train and y_test are vectors containing the classification of each image as a number from 0 to 9. We convert each of these to a format better suited to training with Flux.jl.

[language = Julia] using Flux

# Reshape features for model input x_train = Flux.flatten(x_train) x_test = Flux.flatten(x_test)

# Encode categorical outputs and store y_train = Flux.onehotbatch(y_train, 0:9) y_test = Flux.onehotbatch(y_test, 0:9) data = [(x_train, y_train)]

Features are now stored in a $28^{2}\times N$ Matrix where each column contains pixel data from a single image, and the labels have been converted to a $10\times N$ Flux.OneHotMatrix where each column contains a 1 in the row corresponding to the image’s classification (eg: row 3 for an image showing the number 2) and a 0 otherwise.

We can now construct an LBDN model to train on the MNIST dataset. The larger the model, the better the classification accuracy will be, at the cost of longer training times. The smaller the Lipschitz bound $\gamma$, the more robust the model will be to input perturbations (such as noise in the image). If $\gamma$ is too small, however, it can restrict the model flexibility and limit the achievable performance [Wang+Manchester2023]. For this example, we use a small network of two 64-neuron hidden layers and set a Lipschitz bound of $\gamma=5.0$ just to demonstrate the method.

[language = Julia] using RobustNeuralNetworks

# Model specification nu = 28*28 # Inputs (size of image) ny = 10 # Outputs (classifications) nh = fill(64,2) # Hidden layers γ = 5.0f0 # Lipschitz bound 5.0

# Define parameters,create model model_ps = DenseLBDNParamsT(nu, nh, ny, γ) model = Chain(DiffLBDN(model_ps), Flux.softmax)

The model consists of two parts. The first is a callable DiffLBDN model constructed from its direct parameterization, which is defined by an instance of DenseLBDNParams as per Section 2.3. The output is then converted to a probability distribution using a softmax layer. Note that all AbstractLBDN models can be combined with traditional neural network layers using Flux.Chain.

We could also have constructed the model from a combination of SandwichFC layers. Introduced in [Wang+Manchester2023], the SandwichFC layer is a fully-connected or dense layer with a guaranteed Lipschitz bound of 1.0. We have designed the user interface for SandwichFC to be as similar to that of Flux.Dense as possible. {lstlisting}[language = Julia] model = Chain( (x) -> (sqrt(γ) * x), SandwichFC(nu => nh[1], relu; T), SandwichFC(nh[1] => nh[2], relu; T), (x) -> (sqrt(γ) * x), SandwichFC(nh[2] => ny; output_layer=true, T), Flux.softmax ) This model is equivalent to a dense LBDN constructed with LBDN or DiffLBDN. We have included it as a convenience for users familiar with layer-wise network construction in Flux.jl, and recommend using it interchangeably with DiffLBDN.

A typical loss function for training on datasets with discrete labels is the cross entropy loss. We can use the crossentropy loss function shipped with Flux.jl.

[language = Julia] loss(model,x,y) = Flux.crossentropy(model(x), y)

We train the model over 600 epochs using two learning rates: 1e-3 for the first 300, and 1e-4 for the last 300. We use the Adam optimizer [Kingma+Ba2015] and the default Flux.train! method for convenience. Note that the Flux.train! method updates the learnable parameters each time the model is evaluated on a batch of data, hence our choice of DiffLBDN as a model wrapper.

[language = Julia] # Hyperparameters epochs = 300 lrs = [1e-3,1e-4]

# Train with the Adam optimizer for k in eachindex(lrs) opt_state = Flux.setup(Adam(lrs[k]), model) for i in 1:epochs Flux.train!(loss, model, data, opt_state) end end

Our final model achieves training and test accuracies of approximately 99% and 98%, respectively, as shown in Table 3.1.6. We could improve this further by switching to a convolutional LBDN, as in [Wang+Manchester2023]. Some examples of classifications given by the trained LBDN model are presented in Figure 4.

The main advantage of using an LBDN for image classification is its built-in robustness to noise (or attacks) added to the image. This robustness is a direct benefit of the Lipschitz bound. As outlined in the Section 2.2.3, the Lipschitz bound effectively defines how “smooth” the network is: the smaller the Lipschitz bound, the less the network outputs will change as the inputs vary. For example, small amounts of noise added to the image will be less likely to change its classification. A detailed investigation into this effect is presented in [Wang+Manchester2023].

We can demonstrate the robustness of LBDNs by comparing the model to a standard MLP built from Flux.Dense layers. We first create a dense network with the same layer structure as the LBDN. {lstlisting}[language = Julia] # Initialisation functions init = Flux.glorot_normal initb(n) = Flux.glorot_normal(n)

# Build a dense model dense = Chain( Dense(nu, nh[1], relu; init, bias=initb(nh[1])), Dense(nh[1], nh[2], relu; init, bias=initb(nh[2])), Dense(nh[2], ny; init, bias=initb(ny)), Flux.softmax )

Training the dense model with the same training loop used for the LBDN model results in a model that achieves training and test accuracies of approximately 98% and 97%, respectively, as shown in Table 3.1.6.

As a simple test of robustness, we add uniformly-sampled random noise in the range $[-\epsilon,\epsilon]$ to the pixel data in the test dataset for a range of noise magnitudes $\epsilon\in[0,200/255].$ We record the test accuracy for each perturbation size and store it for plotting. {lstlisting}[language = Julia] using Statistics

# Get test accuracy as we add noise uniform(x) = 2*rand(T, size(x)…) .- 1 compare(y, yh) = maximum(yh, dims=1) .== maximum(y.*yh, dims=1) accuracy(model, x, y) = mean(compare(y, model(x)))

function noisy_test_error(model, ϵ) noisy_xtest = x_test .+ ϵ*uniform(x_test) accuracy(model, noisy_xtest, y_test)*100 end

ϵs = T.(LinRange(0, 200, 10)) ./ 255 lbdn_error = noisy_test_error.((model,), ϵs) dense_error = noisy_test_error.((dense,), ϵs)

Plotting the results in Figure 5 very clearly shows that the dense network, which has no guarantees on its Lipschitz bound, quickly loses its accuracy as small amounts of noise are added to the image. In contrast, the LBDN model maintains its accuracy even when the (maximum) perturbation size is as much as 80% of the maximum pixel values. This is an illustration of why image classification is one of the most promising use-cases for LBDN models. For a more detailed comparison of LBDN with state-of-the-art image classification methods, see [Wang+Manchester2023].

Consider the simple mechanical system shown in Figure 6: a box of mass $m$ sits in a tub of fluid, held between the walls by two springs each with spring constant $k/2.$ The box can be pushed with a force $u.$ Its dynamics are

where $\mu$ is the viscous damping coefficient due to the box moving through the fluid, and $\dot{q},\ddot{q}$ denote the velocity and acceleration of the box, respectively.

We can write this as a (nonlinear) state-space model with state $x=(q,\dot{q})^{\top},$ control input $u,$ and dynamics

This is a continuous-time model of the dynamics. For our purposes, we need a discrete-time model. We can discretize the dynamics using a forward Euler approximation to get

where $\Delta t$ is the time-step. This approximation typically requires a small time-step for numerical stability, but is sufficient for our simple example. If physical accuracy was of concern, one could use a fourth (or higher) order Runge-Kutta scheme.

Our aim is to learn a controller $u=\mathcal{K}_{\theta}(x,q_{\mathrm{ref}}),$ defined by some learnable parameters $\theta,$ that can push the box to any goal position $q_{\mathrm{ref}}$ that we choose. Specifically, we want the box to:

1. 1.

   reach a (stationary) goal position $q_{\mathrm{ref}}$

2. 2.

   within a time period $T$.

The force required to keep the box at a static equilibrium position $q_{\mathrm{ref}}$ is $u_{\mathrm{ref}}=kq_{\mathrm{ref}}$ from Equation 13. We can encode these objectives into a cost function $J_{\theta}$ and write our RL problem as

where $\Delta q_{t}=q_{t}-q_{\mathrm{ref}}$, $\Delta u_{t}=u_{t}-u_{\mathrm{ref}}$, $c_{1},c_{2},c_{3}$ are cost function weights, and the expectation is over different initial and goal positions of the box.

We start by defining the properties of our system and translating the dynamics into Julia code. For this example, we consider a box of mass $m=1$, spring constants $k=5,$ and a viscous damping coefficient $\mu=0.5$. We will simulate the system over $T=4$ s time horizons with a time-step of $\Delta t=0.02$ s.

[language = Julia] m = 1 # Mass (kg) k = 5 # Spring constant (N/m) μ = 0.5 # Viscous damping (kg/m)

Tmax = 4 # Simulation horizon (s) dt = 0.02 # Time step (s) ts = 1:Int(Tmax/dt) # Array of time indices

Now we can generate the training data. Suppose the box always starts at rest from the zero position, and the goal position can be anywhere in the range $q_{\mathrm{ref}}\in[-1,1]$. Our training data consists of a batch of 80 randomly-sampled goal positions and corresponding reference forces $u_{\mathrm{ref}}$. {lstlisting}[language = Julia] nx, nref, batches = 2, 1, 80 x0 = zeros(nx, batches) qref = 2*rand(nref, batches) .- 1 uref = k*qref

It is good practice (and faster) to simulate all simulation batches at once, so we define our dynamics functions to operate on matrices of states and controls. Each row corresponds to a different state or control, and each column corresponds to a simulation for a particular goal position. {lstlisting}[language = Julia] f(x::Matrix,u::Matrix) = [x[2:2,:]; (u[1:1,:] - k*x[1:1,:] - μ*x[2:2,:].^2)/m] fd(x::Matrix,u::Matrix) = x + dt*f(x,u)

RL problems typically involve simulating the system over some time horizon and collecting rewards or costs at each time step. Control policies are trained using approximations of the cost gradient $\nabla J_{\theta}$, as it is often difficult (or impossible) to compute the exact gradient due to the complexity of dynamics simulators. We refer the reader to ReinforcementLearning.jl for further details [Tian++2020].

For this simple example, we can back-propagate directly through the dynamics function fd(x,u). The simulator below takes a batch of initial states, goal positions, and a controller model whose inputs are $[x;q_{\mathrm{ref}}]$. It computes trajectories of states and controls $z=\{[x_{0};u_{0}],\ldots,[x_{T-1};u_{T-1}]\}$. To avoid the well-known issue of array mutation being unsupported in differentiable programming, we use a Zygote.Buffer to iteratively store the outputs [Innes2018b].

[language = Julia] using Zygote: Buffer

function rollout(model, x0, qref) z = Buffer([zero([x0;qref])], length(ts)) x = x0 for t in ts u = model([x;qref]) z[t] = vcat(x,u) x = fd(x,u) end return copy(z) end

After computing these trajectories, we will need a function to evaluate the cost given some weightings $c_{1},c_{2},c_{3}$. {lstlisting}[language = Julia] using Statistics

weights = [10,1,0.1] function _cost(z, qref, uref) Δz = z .- [qref; zero(qref); uref] return mean(sum(weights .* Δz.^2; dims=1)) end cost(z::AbstractVector, qref, uref) = mean(_cost.(z, (qref,), (uref,)))

We will train an LBDN controller with a Lipschitz bound of $\gamma=20$. Its inputs are the state $x_{t}$ and goal position $q_{\mathrm{ref}}$, while its outputs are the control force $u_{t}$. We have chosen a model with two hidden layers each of 32 neurons just as an example. For details on how Lipschitz bounds can be useful in learning robust controllers, please see [Barbara++2023].

[language = Julia] using RobustNeuralNetworks

T = Float64 γ = 20 # Lipschitz bound nu = nx + nref # Inputs (x and reference) ny = 1 # Outputs (control action) nh = fill(32, 2) # Hidden layers model_ps = DenseLBDNParamsT(nu, nh, ny, γ)

In constructing a loss function for this problem, we refer to Section 2.3.3. The model_ps contain all information required to define a dense LBDN model. However, model_ps is not a model that can be evaluated on data: it is a model parameterization, and contains the learnable parameters $\theta$. To train an LBDN given some data, we construct the model within the loss function using the LBDN wrapper so that the mapping from direct to explicit parameters is captured during back-propagation. Our loss function therefore includes the following three components.

[language = Julia] function loss(model_ps, x0, qref, uref) model = LBDN(model_ps) # Model z = rollout(model, x0, qref) # Simulation return cost(z, qref, uref) # Cost end

Having set up the RL problem, all that remains is to train the controller. The function below trains a model and keeps track of the training loss tloss (cost $J_{\theta}$) for each simulation in our batch of 80. Training is performed with the Adam optimizer over 250 epochs with a learning rate of $10^{-3}$.

[language = Julia] using Flux

function train_box_ctrl!( model_ps, loss_func; epochs=250, lr=1e-3 ) costs = VectorFloat64() opt_state = Flux.setup(Adam(lr), model_ps) for k in 1:epochs

tloss, dJ = Flux.withgradient( loss_func, model