Modeling Non Linear Friction Model using UDEs

Friction between moving bodies is not trivial to model. There have been idealised linear models which are not always useful in complicated systems. There have been many theories and non linear models which we can use, but they are not perfect. The aim of this tutorial to use Universal Differential Equations to showcase how we can embed a neural network to learn an unknown non linear friction model.

Julia Environment

First, lets import the required packages.

using ModelingToolkitNeuralNets
using ModelingToolkit
import ModelingToolkit.t_nounits as t
import ModelingToolkit.D_nounits as Dt
using ModelingToolkitStandardLibrary.Blocks
using OrdinaryDiffEq
using Optimization
using OptimizationOptimisers: Adam
using SciMLStructures
using SciMLStructures: Tunable
using SymbolicIndexingInterface
using StableRNGs
using Lux
using Plots

Problem Setup

Let's use the friction model presented in https://www.mathworks.com/help/simscape/ref/translationalfriction.html for generating data.

Fbrk = 100.0
vbrk = 10.0
Fc = 80.0
vst = vbrk / 10
vcol = vbrk * sqrt(2)
function friction(v)
    sqrt(2 * MathConstants.e) * (Fbrk - Fc) * exp(-(v / vst)^2) * (v / vst) +
    Fc * tanh(v / vcol)
end

friction (generic function with 1 method)

Next, we define the model - an object sliding in 1D plane with a constant force Fu acting on it and friction force opposing the motion.

function friction_true()
    @variables y(t) = 0.0
    @constants Fu = 120.0
    eqs = [
        Dt(y) ~ Fu - friction(y)
    ]
    return ODESystem(eqs, t, name = :friction_true)
end

friction_true (generic function with 1 method)

Now that we have defined the model, we will simulate it from 0 to 0.1 seconds.

model_true = structural_simplify(friction_true())
prob_true = ODEProblem(model_true, [], (0, 0.1), [])
sol_ref = solve(prob_true, Rodas4(); saveat = 0.001)

retcode: Success
Interpolation: 1st order linear
t: 101-element Vector{Float64}:
 0.0
 0.001
 0.002
 0.003
 0.004
 0.005
 0.006
 0.007
 0.008
 0.009
 ⋮
 0.092
 0.093
 0.094
 0.095
 0.096
 0.097
 0.098
 0.099
 0.1
u: 101-element Vector{Vector{Float64}}:
 [0.0]
 [0.11693525634386417]
 [0.22815063944906144]
 [0.3344560537929039]
 [0.4368038331727264]
 [0.5361810711751653]
 [0.6335357754033413]
 [0.7297484315722541]
 [0.8255832780234691]
 [0.9216604041698457]
 ⋮
 [8.58644383781187]
 [8.662919111645309]
 [8.739090192388025]
 [8.814960147942514]
 [8.890532046211273]
 [8.965808955096799]
 [9.040793942501585]
 [9.115490076328129]
 [9.189900424478925]

Let's plot it.

scatter(sol_ref, label = "velocity")

That was the velocity. Let's also plot the friction force acting on the object throughout the simulation.

scatter(sol_ref.t, friction.(first.(sol_ref.u)), label = "friction force")

Model Setup

Now, we will try to learn the same friction model using a neural network. We will use NeuralNetworkBlock to define neural network as a component. The input of the neural network is the velocity and the output is the friction force. We connect the neural network with the model using RealInputArray and RealOutputArray blocks.

function friction_ude(Fu)
    @variables y(t) = 0.0
    @constants Fu = Fu
    @named nn_in = RealInputArray(nin = 1)
    @named nn_out = RealOutputArray(nout = 1)
    eqs = [Dt(y) ~ Fu - nn_in.u[1]
           y ~ nn_out.u[1]]
    return ODESystem(eqs, t, name = :friction, systems = [nn_in, nn_out])
end

Fu = 120.0
model = friction_ude(Fu)

chain = Lux.Chain(
    Lux.Dense(1 => 10, Lux.mish, use_bias = false),
    Lux.Dense(10 => 10, Lux.mish, use_bias = false),
    Lux.Dense(10 => 1, use_bias = false)
)
@named nn = NeuralNetworkBlock(1, 1; chain = chain, rng = StableRNG(1111))

eqs = [connect(model.nn_in, nn.output)
       connect(model.nn_out, nn.input)]

ude_sys = complete(ODESystem(eqs, t, systems = [model, nn], name = :ude_sys))
sys = structural_simplify(ude_sys)

\[ \begin{align} \frac{\mathrm{d} \mathtt{friction.y}\left( t \right)}{\mathrm{d}t} &= \mathtt{friction.Fu} - \mathtt{friction.nn\_in.u}\left( t \right)_{1} \\ 0 &= - \mathtt{friction.nn\_out.u}\left( t \right)_{1} + \mathtt{nn.input.u}\left( t \right)_{1} \end{align} \]

Optimization Setup

We now setup the loss function and the optimization loop.

function loss(x, (prob, sol_ref, get_vars, get_refs, set_x))
    new_p = set_x(prob, x)
    new_prob = remake(prob, p = new_p, u0 = eltype(x).(prob.u0))
    ts = sol_ref.t
    new_sol = solve(new_prob, Rodas4(), saveat = ts, abstol = 1e-8, reltol = 1e-8)
    loss = zero(eltype(x))
    for i in eachindex(new_sol.u)
        loss += sum(abs2.(get_vars(new_sol, i) .- get_refs(sol_ref, i)))
    end
    if SciMLBase.successful_retcode(new_sol)
        loss
    else
        Inf
    end
end

of = OptimizationFunction{true}(loss, AutoForwardDiff())

prob = ODEProblem(sys, [], (0, 0.1), [])
get_vars = getu(sys, [sys.friction.y])
get_refs = getu(model_true, [model_true.y])
set_x = setp_oop(sys, sys.nn.p)
x0 = default_values(sys)[sys.nn.p]

cb = (opt_state, loss) -> begin
    @info "step $(opt_state.iter), loss: $loss"
    return false
end

op = OptimizationProblem(of, x0, (prob, sol_ref, get_vars, get_refs, set_x))
res = solve(op, Adam(5e-3); maxiters = 10000, callback = cb)

retcode: Default
u: 120-element Vector{Float64}:
 -2.363635617199658
 -2.4192743321986514
  2.055321602454787
 -0.2999252069237061
  0.9657348368044494
 -2.4086786500487363
 -2.321253698035093
  0.5512956242140039
 -2.4360068303020124
 -2.425651887694722
  ⋮
 -2.1175791468725773
  1.231137921644243
 -1.046475488871763
 -0.25115604088590265
 -1.5963390505289439
  1.4534888331289002
 -2.282535867419842
 -1.4113655162853294
 -1.5278017288067132

Visualization of results

We now have a trained neural network! We can check whether running the simulation of the model embedded with the neural network matches the data or not.

res_p = set_x(prob, res.u)
res_prob = remake(prob, p = res_p)
res_sol = solve(res_prob, Rodas4(), saveat = sol_ref.t)

Also, it would be interesting to check the simulation before the training to get an idea of the starting point of the network.

initial_sol = solve(prob, Rodas4(), saveat = sol_ref.t)

retcode: Success
Interpolation: 1st order linear
t: 101-element Vector{Float64}:
 0.0
 0.001
 0.002
 0.003
 0.004
 0.005
 0.006
 0.007
 0.008
 0.009
 ⋮
 0.092
 0.093
 0.094
 0.095
 0.096
 0.097
 0.098
 0.099
 0.1
u: 101-element Vector{Vector{Float64}}:
 [0.0, 0.0]
 [0.12001891802081038, 0.12001891802081037]
 [0.24007647700875054, 0.24007647700875048]
 [0.3601739747958169, 0.3601739747958169]
 [0.4803131095528584, 0.48031310955285844]
 [0.6004927757829802, 0.6004927757829802]
 [0.7207126193353081, 0.7207126193353081]
 [0.8409730642995663, 0.8409730642995663]
 [0.9612712084090433, 0.9612712084090433]
 [1.0816044728123086, 1.0816044728123086]
 ⋮
 [11.145787875099332, 11.145787875099332]
 [11.267953036754593, 11.267953036754593]
 [11.390141551847606, 11.390141551847606]
 [11.512353442656476, 11.512353442656476]
 [11.634588731459306, 11.634588731459306]
 [11.756847440534198, 11.756847440534198]
 [11.879129592159257, 11.879129592159257]
 [12.001435208612586, 12.001435208612586]
 [12.123764312172286, 12.123764312172286]

Now we plot it.

scatter(sol_ref, idxs = [model_true.y], label = "ground truth velocity")
plot!(res_sol, idxs = [sys.friction.y], label = "velocity after training")
plot!(initial_sol, idxs = [sys.friction.y], label = "velocity before training")

It matches the data well! Let's also check the predictions for the friction force and whether the network learnt the friction model or not.

scatter(sol_ref.t, friction.(first.(sol_ref.u)), label = "ground truth friction")
plot!(res_sol.t, getindex.(res_sol[sys.nn.output.u], 1),
    label = "friction from neural network")

It learns the friction model well!