Merge pull request #107 from acfr/bugfix/box-dynamics

Bugfix: box dynamics

docs/src/examples/box_obsv.md

@@ -47,7 +47,8 @@ k = 5                   # Spring constant (N/m)
 nx = 2                  # Number of states
 
 # Continuous and discrete dynamics and measurements
-f(x::Matrix,u::Matrix) = [x[2:2,:]; (u[1:1,:] - k*x[1:1,:] - μ*x[2:2,:].^2)/m]
+_visc(v::Matrix) = μ * v .* abs.(v)
+f(x::Matrix,u::Matrix) = [x[2:2,:]; (u[1:1,:] - k*x[1:1,:] - _visc(x[2:2,:]))/m]
 fd(x,u) = x + dt*f(x,u)
 gd(x::Matrix) = x[1:1,:]
 ```

docs/src/examples/lbdn_mnist.md

@@ -132,7 +132,7 @@ function train_mnist!(model, data; num_epochs=300, lrs=[1e-3,1e-4])
 end
 
 # Train and save the model for later
-train_mnist!(model, data)
+train_mnist!(model, train_data)
 bson("lbdn_mnist.bson", Dict("model" => model))
 ```
 

docs/src/examples/rl.md

@@ -8,7 +8,7 @@ In this example, we'll demonstrate how to train an LBDN controller with *Reinfor
 
 Let's consider the simple mechanical system shown below: a box of mass ``m`` sits in a tub of fluid, held between the walls of the tub by two springs, each with spring constant ``k/2.`` We can push the box with force ``u.`` Its dynamics are
 ```math
-m\ddot{q} = u - kq - \mu \dot{q}^2
+m\ddot{q} = u - kq - \mu \dot{q} |\dot{q}|
 ```
 where ``\mu`` is the viscous friction coefficient due to the box moving through the fluid.
 
@@ -19,7 +19,7 @@ where ``\mu`` is the viscous friction coefficient due to the box moving through
 We can write this as a (nonlinear) state-space model with state ``x = (q,\dot{q})^\top,`` control input ``u,`` and dynamics
 ```math
 \dot{x} = f(x,u) = \begin{bmatrix}
-\dot{q} \\ (u - kq - \mu \dot{q}^2)/m
+\dot{q} \\ (u - kq - \mu \dot{q} |\dot{q}|)/m
 \end{bmatrix}.
 ```
 This is a continous-time model of the dynamics. For our purposes, we'll need a discrete-time model. We can discretise the dynamics using a [forward Euler approximation](https://en.wikipedia.org/wiki/Euler_method) such that
@@ -68,7 +68,8 @@ uref = k*qref
 
 It's good practice (and faster) to simulate all of these simulation batches at once, so we define our dynamics functions to operate on matrices of states and controls. Each row is a different state or control, and each column corresponds to a simulation for a particular goal position.
 ```julia
-f(x::Matrix,u::Matrix) = [x[2:2,:]; (u[1:1,:] - k*x[1:1,:] - μ*x[2:2,:].^2)/m]
+_visc(v::Matrix) = μ * v .* abs.(v)
+f(x::Matrix,u::Matrix) = [x[2:2,:]; (u[1:1,:] - k*x[1:1,:] - _visc(x[2:2,:]))/m]
 fd(x::Matrix,u::Matrix) = x + dt*f(x,u)
 ```
 

examples/src/lbdn_mnist.jl

@@ -79,8 +79,7 @@ end
 
 # Train and save the model for later use
 train_mnist!(model, train_data)
-stop_here
-# bson("assets/lbdn-mnist/lbdn_mnist.bson", Dict("model" => model))
+bson("assets/lbdn-mnist/lbdn_mnist.bson", Dict("model" => model))
 model = BSON.load("assets/lbdn-mnist/lbdn_mnist.bson")["model"]
 
 # Print final results

examples/src/lbdn_rl.jl

@@ -35,8 +35,9 @@ x0 = zeros(nx, batches)
 qref = 2*rand(rng, nref, batches) .- 1
 uref = k*qref
 
-# Continuous  and discrete dynamics
-f(x::Matrix,u::Matrix) = [x[2:2,:]; (u[1:1,:] - k*x[1:1,:] - μ*x[2:2,:].^2)/m]
+# Continuous and discrete dynamics
+_visc(v::Matrix) = μ * v .* abs.(v)
+f(x::Matrix,u::Matrix) = [x[2:2,:]; (u[1:1,:] - k*x[1:1,:] - _visc(x[2:2,:]))/m]
 fd(x::Matrix,u::Matrix) = x + dt*f(x,u)
 
 # Simulate the system given initial condition and a controller
@@ -97,7 +98,7 @@ function train_box_ctrl!(model_ps, loss_func; lr=1e-3, epochs=250, verbose=false
     return costs
 end
 
-costs = train_box_ctrl!(model_ps, loss)
+costs = train_box_ctrl!(model_ps, loss; verbose=true)
 
 
 # -------------------------
@@ -126,8 +127,8 @@ function plot_box_learning(costs, z, qr)
     ga = f1[1,1] = GridLayout()
 
     ax0 = Axis(ga[1,1], xlabel="Training epochs", ylabel="Cost")
-    ax1 = Axis(ga[1,2], xlabel="Time (s))", ylabel="Position error (m)", )
-    ax2 = Axis(ga[2,1], xlabel="Time (s))", ylabel="Velocity (m/s)")
+    ax1 = Axis(ga[1,2], xlabel="Time (s)", ylabel="Position error (m)", )
+    ax2 = Axis(ga[2,1], xlabel="Time (s)", ylabel="Velocity (m/s)")
     ax3 = Axis(ga[2,2], xlabel="Time (s)", ylabel="Control error (N)")
 
     lines!(ax0, costs, color=:black)
@@ -192,4 +193,3 @@ xlims!(ax, [sizes[1], sizes[end]])
 axislegend(ax, position=:lt)
 display(f1)
 save("../results/lbdn_rl_comptime.svg", f1)
-

examples/src/ren_obsv_box.jl

@@ -24,7 +24,8 @@ k = 5                   # Spring constant (N/m)
 nx = 2                  # Number of states
 
 # Continuous and discrete dynamics and measurements
-f(x::Matrix,u::Matrix) = [x[2:2,:]; (u[1:1,:] - k*x[1:1,:] - μ*x[2:2,:].^2)/m]
+_visc(v::Matrix) = μ * v .* abs.(v)
+f(x::Matrix,u::Matrix) = [x[2:2,:]; (u[1:1,:] - k*x[1:1,:] - _visc(x[2:2,:]))/m]
 fd(x,u) = x + dt*f(x,u)
 gd(x::Matrix) = x[1:1,:]
 
@@ -56,10 +57,10 @@ data = zip(Xn[indx], Xt[indx], observer_data[indx])
 # Train a model
 
 # Define a REN model for the observer
-nv = 100
+nv = 200
 nu = size(observer_data[1], 1)
 ny = nx
-model_ps = ContractingRENParams{Float64}(nu, nx, nv, ny; output_map=false, rng)
+model_ps = ContractingRENParams{Float32}(nu, nx, nv, ny; output_map=false, rng)
 model = DiffREN(model_ps)
 
 # Loss function: one step ahead error (average over time)
@@ -69,7 +70,7 @@ function loss(model, xn, xt, inputs)
 end
 
 # Train the model
-function train_observer!(model, data; epochs=100, lr=1e-3, min_lr=1e-4)
+function train_observer!(model, data; epochs=50, lr=1e-3, min_lr=1e-6)
 
     opt_state = Flux.setup(Adam(lr), model)
     mean_loss = [1e5]
@@ -85,7 +86,7 @@ function train_observer!(model, data; epochs=100, lr=1e-3, min_lr=1e-4)
 
         # Drop learning rate if mean loss is stuck or growing
         push!(mean_loss, mean(batch_loss))
-        if (mean_loss[end] >= mean_loss[end-1]) && !(lr <= min_lr)
+        if (mean_loss[end] >= mean_loss[end-1]) && !(lr < min_lr || lr ≈ min_lr)
             lr = 0.1lr
             Flux.adjust!(opt_state, lr)
         end
@@ -100,7 +101,7 @@ tloss = train_observer!(model, data)
 
 # Generate test data (a bunch of initial conditions)
 batches   = 50
-ts_test   = 1:Int(10/dt)
+ts_test   = 1:Int(20/dt)
 u_test    = fill(zeros(1, batches), length(ts_test))
 x_test    = fill(zeros(nx,batches), length(ts_test))
 x_test[1] = 0.2*(2*rand(rng, nx, batches) .-1)

test/ParameterTypes/contraction.jl

@@ -6,6 +6,8 @@ using Test
 
 # include("../test_utils.jl")
 
+rng = MersenneTwister(42)
+
 """
 Test that the contracting REN actually does contract.
 
@@ -17,15 +19,15 @@ nu, nx, nv, ny = 4, 5, 10, 5
 
 ren_ps = ContractingRENParams{Float64}(
     nu, nx, nv, ny; 
-    init=:cholesky, αbar=ᾱ, polar_param=false, output_map=false
+    init=:cholesky, αbar=ᾱ, polar_param=false, output_map=false, rng
 )
 ren = REN(ren_ps)
 
 # Same inputs. different initial conditions
-u0 = randn(nu, batches)
+u0 = randn(rng, nu, batches)
 
-x0 = randn(nx, batches)
-x1 = randn(nx, batches)
+x0 = randn(rng, nx, batches)
+x1 = randn(rng, nx, batches)
 
 # Simulate
 x0n, y0 = ren(x0, u0)

test/ParameterTypes/dense_lbdn.jl

@@ -6,6 +6,8 @@ using Test
 
 # include("../test_utils.jl")
 
+rng = MersenneTwister(42)
+
 """
 Test that the model satisfies a specified Lipschitz bound
 """
@@ -14,12 +16,12 @@ nu, ny = 4, 2
 nh = [10, 5, 20, 4]
 γ = 1e-5
 
-lbdn_ps = DenseLBDNParams{Float64}(nu, nh, ny, γ)
+lbdn_ps = DenseLBDNParams{Float64}(nu, nh, ny, γ; rng)
 lbdn = LBDN(lbdn_ps)
 
 # Different inputs with different initial conditions
-u0 = randn(nu, batches)
-u1 = u0 .+ 0.001*rand(nu, batches)
+u0 = randn(rng, nu, batches)
+u1 = u0 .+ 0.001*rand(rng, nu, batches)
 
 # Simulate
 y0 = lbdn(u0)

test/ParameterTypes/general_behavioural_constrains.jl

@@ -7,29 +7,31 @@ using Test
 
 # include("../test_utils.jl")
 
+rng = MersenneTwister(42)
+
 """
 Test that the behavioural constraints are satisfied
 """
 batches = 42
 nu, nx, nv, ny = 10, 5, 10, 20
 
 # Generate random matrices
-X = randn(ny,ny)
-Y = randn(nu,nu)
-S = rand(nu,ny)
+X = randn(rng, ny,ny)
+Y = randn(rng, nu,nu)
+S = rand(rng, nu,ny)
 
 Q = -X'*X
 R = S * (Q \ S') + Y'*Y
 
-ren_ps = GeneralRENParams{Float64}(nu, nx, nv, ny, Q, S, R)
+ren_ps = GeneralRENParams{Float64}(nu, nx, nv, ny, Q, S, R; rng)
 ren = REN(ren_ps)
 
 # Different inputs with different initial conditions
-u0 = 10*randn(nu, batches)
-u1 = rand(nu, batches)
+u0 = 10*randn(rng, nu, batches)
+u1 = rand(rng, nu, batches)
 
-x0 = randn(nx, batches)
-x1 = randn(nx, batches)
+x0 = randn(rng, nx, batches)
+x1 = randn(rng, nx, batches)
 
 # Simulate
 x0n, y0 = ren(x0, u0)

test/ParameterTypes/lipschitz_bound.jl

@@ -6,22 +6,24 @@ using Test
 
 # include("../test_utils.jl")
 
+rng = MersenneTwister(42)
+
 """
 Test that the model satisfies a specified Lipschitz bound
 """
 batches = 42
 nu, nx, nv, ny = 4, 5, 10, 2
 γ = 10
 
-ren_ps = LipschitzRENParams{Float64}(nu, nx, nv, ny, γ)
+ren_ps = LipschitzRENParams{Float64}(nu, nx, nv, ny, γ; rng)
 ren = REN(ren_ps)
 
 # Different inputs with different initial conditions
-u0 = 10*randn(nu, batches)
-u1 = rand(nu, batches)
+u0 = 10*randn(rng, nu, batches)
+u1 = rand(rng, nu, batches)
 
-x0 = randn(nx, batches)
-x1 = randn(nx, batches)
+x0 = randn(rng, nx, batches)
+x1 = randn(rng, nx, batches)
 
 # Simulate
 x0n, y0 = ren(x0, u0)

test/ParameterTypes/passivity.jl

@@ -7,6 +7,8 @@ using Test
 
 # include("../test_utils.jl")
 
+rng = MersenneTwister(42)
+
 """
 Test passivity inequality
 """
@@ -15,15 +17,15 @@ nu, nx, nv, ny = 6, 5, 10, 6
 T = 100
 
 # Test constructors
-ren_ps = PassiveRENParams{Float64}(nu, nx, nv, ny; init=:random, ν= 1.0)
+ren_ps = PassiveRENParams{Float64}(nu, nx, nv, ny; init=:random, ν= 1.0, rng)
 ren = REN(ren_ps)
 
 # Different inputs with different initial conditions
-u0 = 10*randn(nu, batches)
-u1 = rand(nu, batches)
+u0 = 10*randn(rng, nu, batches)
+u1 = rand(rng, nu, batches)
 
-x0 = randn(nx, batches)
-x1 = randn(nx, batches)
+x0 = randn(rng, nx, batches)
+x1 = randn(rng, nx, batches)
 
 # Simulate
 x0n, y0 = ren(x0, u0)

test/Wrappers/diff_lbdn.jl

@@ -5,18 +5,20 @@ using Random
 using RobustNeuralNetworks
 using Test
 
+rng = MersenneTwister(42)
+
 """
 Test that backpropagation runs and parameters change
 """
 batches = 10
 nu, ny, γ = 2, 3, 1
 nh = [10,5]
-model_ps = DenseLBDNParams{Float32}(nu, nh, ny, γ; learn_γ=true)
+model_ps = DenseLBDNParams{Float32}(nu, nh, ny, γ; learn_γ=true, rng)
 model = DiffLBDN(model_ps)
 
 # Dummy data
-us = randn(nu, batches)
-ys = randn(ny, batches)
+us = randn(rng, nu, batches)
+ys = randn(rng, ny, batches)
 data = [(us[:,k], ys[:,k]) for k in 1:batches]
 
 # Dummy loss function just for testing

test/Wrappers/diff_ren.jl

@@ -5,18 +5,20 @@ using Random
 using RobustNeuralNetworks
 using Test
 
+rng = MersenneTwister(42)
+
 """
 Test that backpropagation runs and parameters change
 """
 batches = 10
 nu, nx, nv, ny = 4, 5, 10, 2
 γ = 10
-ren_ps = LipschitzRENParams{Float64}(nu, nx, nv, ny, γ)
+ren_ps = LipschitzRENParams{Float64}(nu, nx, nv, ny, γ; rng)
 model = DiffREN(ren_ps)
 
 # Dummy data
-us = randn(nu, batches)
-ys = randn(ny, batches)
+us = randn(rng, nu, batches)
+ys = randn(rng, ny, batches)
 data = [(us[:,k], ys[:,k]) for k in 1:batches]
 
 # Dummy loss function just for testing

test/Wrappers/wrap_ren.jl

@@ -5,6 +5,8 @@ using Random
 using RobustNeuralNetworks
 using Test
 
+rng = MersenneTwister(42)
+
 """
 Test REN wrapper with General REN params
 """
@@ -15,15 +17,15 @@ Q = Matrix{Float64}(-I(ny))
 R = 0.1^2 * Matrix{Float64}(I(nu))
 S = zeros(Float64, nu, ny)
 
-ren_ps = GeneralRENParams{Float64}(nu, nx, nv, ny, Q, S, R)
+ren_ps = GeneralRENParams{Float64}(nu, nx, nv, ny, Q, S, R; rng)
 ren1 = WrapREN(ren_ps)
 
 x0 = init_states(ren1, batches)
-u0 = randn(nu, batches)
+u0 = randn(rng, nu, batches)
 
 # Update the model after changing a parameter
 old_B2 = deepcopy(ren1.explicit.B2)
-ren1.params.direct.B2 .*= rand(size(ren1.params.direct.B2)...)
+ren1.params.direct.B2 .*= rand(rng, size(ren1.params.direct.B2)...)
 
 x1, y1 = ren1(x0, u0)
 update_explicit!(ren1)

test/Wrappers/zero_dim.jl

@@ -5,17 +5,19 @@ using Random
 using RobustNeuralNetworks
 using Test
 
+rng = MersenneTwister(42)
+
 """
 Test that backpropagation runs when nx = 0 and nv = 0
 """
 batches = 10
 nu, nx, nv, ny = 4, 0, 0, 2
 γ = 10
-model_ps = LipschitzRENParams{Float64}(nu, nx, nv, ny, γ)
+model_ps = LipschitzRENParams{Float64}(nu, nx, nv, ny, γ; rng)
 
 # Dummy data
-us = randn(nu, batches)
-ys = randn(ny, batches)
+us = randn(rng, nu, batches)
+ys = randn(rng, ny, batches)
 data = [(us, ys)]
 
 # Dummy loss function just for testing