Testing FastDifferentiation.jl for TORA.jl

[ferrolho]

Hi, Brian! 👋 I am currently using ForwardDiff.jl within TORA.jl for computing Jacobians of user-defined vector-valued functions. However, I noticed that FastDifferentiation.jl (this package) may provide quicker Jacobian calculations compared to ForwardDiff.jl, and so I would like to give it a go.

I've created a Jupyter notebook, TestFastDifferentiation.ipynb, on a new branch (hf/test-FastDifferentiation.jl), where I want to benchmark ForwardDiff.jl and FastDifferentiation.jl side-by-side. I've managed to create the variables with

x = make_variables(:x, robot.n_q + robot.n_v + robot.n_τ + robot.n_q + robot.n_v)

but I can't figure out the correct syntax for the actual evaluation of the Jacobian.

P.s. the function I am trying to differentiate is forward_dynamics_defects!.

[brianguenter]

Is this what you intended:

function test_robot()
    robot = TORA.create_robot_franka("panda_arm", vis)
    x = make_variables(:x, robot.n_q + robot.n_v + robot.n_τ + robot.n_q + robot.n_v)
    defects = zeros(robot.n_q + robot.n_v)
    dyn = TORA.forward_dynamics_defects!(defects, robot, x, dt)
    jac_defects = jacobian(dyn, x)
    dyn_func = make_function(jac_defects, x)
    dyn_func(rand(robot.n_q + robot.n_v + robot.n_τ + robot.n_q + robot.n_v))
end

If this looks correct I'll get this working. There is currently a type promotion error in the muladd function that might take a little while to fix.

[ferrolho]

Yes, this looks correct. Thanks!

I am not sure if the muladd issue you are referring to is the same as this one, but I get

MethodError: no method matching AbstractFloat(::FastDifferentiation.Node{Symbol, 0})

at the line

dyn = TORA.forward_dynamics_defects!(defects, robot, x, dt)

[brianguenter]

Yeah, I think that is a separate error related to sincos. This function wasn't defined for Node types because SymbolicUtils.number_methods doesn't define it and number methods is what I use to define all the functions for Node types. I've fixed it in my local code, not sure if that has been pushed to the repo yet. But the promotion error is still there. Have to attend to other things for a few hours but should have this fixed by tonight.

[ferrolho]

The issue has been closed, but I just tested this on the latest main and there is this error:

promotion of types FastDifferentiation.Node{Int64, 0}, FastDifferentiation.Node{Float64, 0} and FastDifferentiation.Node{Int64, 0} failed to change any arguments

[brianguenter]

sorry about that I must have accidentally closed this. I'm still working on the type promotion error.

[brianguenter]

Fixed the type promotion error, which exposed problems in the LinearAlgebra code.

Your code calls dynamics_solve! which calls LU which has conditionals which FastDifferentiation.jl doesn't yet support.

Symbolics.jl has an LU factorization which uses a heuristic to do partial pivoting since the value of the matrix elements is not known at the time the symbolic expression is created. Instead of comparing numeric values they compare the number of terms in the symbolic expressions. No idea how well this heuristic works or whether this would result in good pivot choices.

I can probably implement this symbolic LU in FastDifferentiation. Would that be useful?

[brianguenter]

To make this work the signature of the symbolic LU would have to exactly match the signature of the LU you are calling, except for the element type of the arrays. I'm reasonably confident this will be simple.

[ferrolho]

Would that be useful?

It would be nice if FastDifferentiation.jl could provide faster Jacobian evaluation for methods that use RigidBodyDynamics.jl; but please don't take this as a request. Of course, it would be useful, but you should only work on it if you'd like to do so. 😅

To make this work the signature of the symbolic LU would have to exactly match the signature of the LU you are calling, except for the element type of the arrays. I'm reasonably confident this will be simple.

Would this require changing the code in dynamics_solve! (from https://github.com/JuliaRobotics/RigidBodyDynamics.jl)?

[brianguenter]

I think it should be possible to overload the \ operator specialized to matrices of Node type. Then it would be completely transparent.

The partial pivoting would be done symbolically in the LU factorization rather than based on actual numeric values of matrix elements ,which might in some cases lead to slightly worse condition numbers than you would expect. Not sure if this will be a big problem.

Other people have needed LU factorization as well so it's worth doing.

[baggepinnen]

The \ matrix operator appears in MTK systems compiled with JuliaSimCompiler as well, differentiating through them efficiently has proven rather difficult (especially if they represent multibody systems)

[brianguenter]

I think I have figured out how to make FastDifferentiation work with ChainRules. I'm currently working on conditionals, so adding ChainRules support will come once I complete conditionals. ChainRules is tricky because FastDifferentiation is inherently a scalar system now and ChainRules is all about Jacobian vector products, which is not really a thing in FastDifferentiation.