hessians_rbd.jl

old = false

Pkg.add("RigidBodyDynamics")
Pkg.checkout("RigidBodyDynamics")
Pkg.add("ForwardDiff")
if !old
  Pkg.checkout("ForwardDiff") # want https://github.com/JuliaDiff/ForwardDiff.jl/pull/252
else
  Pkg.pin("ForwardDiff", v"0.4.2")
end


using RigidBodyDynamics
using JuMP
using Ipopt
using StaticArrays

# define some mechanism, taken from RigidBodyDynamics examples
g = -9.81 # gravitational acceleration in z-direction
world = RigidBody{Float64}("world")
mechanism = Mechanism(world; gravity = SVector(0, 0, g))

axis = SVector(0., 1., 0.) # joint axis
I_1 = 0.333 # moment of inertia about joint axis
c_1 = -0.5 # center of mass location with respect to joint axis
m_1 = 1. # mass
frame1 = CartesianFrame3D("upper_link") # the reference frame in which the spatial inertia will be expressed
inertia1 = SpatialInertia(frame1, I_1 * axis * axis.', m_1 * SVector(0, 0, c_1), m_1)

upperlink = RigidBody(inertia1)
shoulder = Joint("shoulder", Revolute(axis))
before_shoulder_to_world = eye(Transform3D, frame_before(shoulder), default_frame(world))
attach!(mechanism, world, upperlink, shoulder, joint_pose = before_shoulder_to_world)

l_1 = -1. # length of the upper link
I_2 = 0.333 # moment of inertia about joint axis
c_2 = -0.5 # center of mass location with respect to joint axis
m_2 = 1. # mass
inertia2 = SpatialInertia(CartesianFrame3D("lower_link"), I_2 * axis * axis.', m_2 * SVector(0, 0, c_2), m_2)
lowerlink = RigidBody(inertia2)
elbow = Joint("elbow", Revolute(axis))
before_elbow_to_after_shoulder = Transform3D(frame_before(elbow), frame_after(shoulder), SVector(0, 0, l_1))
attach!(mechanism, upperlink, lowerlink, elbow, joint_pose = before_elbow_to_after_shoulder)


import Base.copy!
function copy!(dest::MechanismState, src::MechanismState)
    set_configuration!(dest, configuration(src))
    set_velocity!(dest, velocity(src))
end

function dynamics_implicit_coulomb!(vd::AbstractArray, sd::AbstractArray, t, state, τ_coulomb_friction, dynamics_result_workspace)
    damping = 0.0
    spring = 10.0
    τ = - spring * configuration(state) + -damping * velocity(state)
    dynamics!(dynamics_result_workspace, state, τ_coulomb_friction + τ)
    copy!(vd, dynamics_result_workspace.v̇)
    copy!(sd, dynamics_result_workspace.ṡ)
    nothing
end

#using RigidBodyDynamics.OdeIntegrators
function evaluate_friction_kinetic_energy{T}(state, t, Δt, τ_tentative_friction::AbstractArray{T})
    # should not mutate `state`

    dynamics_result_workspace = DynamicsResult{T}(mechanism)
    function dynamics!(vd::AbstractArray, sd::AbstractArray, t, state)
        dynamics_implicit_coulomb!(vd::AbstractArray, sd::AbstractArray, t, state,
            τ_tentative_friction, dynamics_result_workspace)
    end

    # guess at a large-enough ExpandingStorage
    integrator = RigidBodyDynamics.OdeIntegrators.MuntheKaasIntegrator(dynamics!,
      RigidBodyDynamics.OdeIntegrators.runge_kutta_4(T),
      RigidBodyDynamics.OdeIntegrators.ExpandingStorage{T}(100))

    state_copy = MechanismState{T}(state.mechanism)

    # dummy call to allocate Arrays for stages.
    RigidBodyDynamics.OdeIntegrators.integrate(integrator, state_copy, 0.0, 1e-3)
    copy!(state_copy, state)
    RigidBodyDynamics.OdeIntegrators.step(integrator, t, state_copy, Δt)
    return kinetic_energy(state_copy, state.mechanism.graph.vertices)
end

# demonstrate (tagged) Dual numbers through evaluate_friction_kinetic_energy
state = MechanismState{Float64}(mechanism)
zero!(state)
set_configuration!(state, [20.0, 0.0])
@show configuration(state)
@show velocity(state)

if old
  τ = fill(ForwardDiff.Dual(10.0, 1.0), num_velocities(mechanism))
  τ2 = fill(Dual(Dual(1.0, 2.0), Dual(1.0,2.0)), num_velocities(mechanism))
else
  # just show that tags don't introduce problems.
  τ = fill(ForwardDiff.Dual{:some_tag}(10.0, 1.0), num_velocities(mechanism))
  τ2 = fill(Dual{:a}(Dual{:b}(1.0, 2.0), Dual{:b}(1.0,2.0)), num_velocities(mechanism))
end

println("un-nested dual numbers")
@show evaluate_friction_kinetic_energy(state, 0.0, 1e-3, τ)
println("nested dual numbers")
@show evaluate_friction_kinetic_energy(state, 0.0, 1e-3, τ2)

f = (τ) -> evaluate_friction_kinetic_energy(state, 0.0, 1e-3, τ)
@show ForwardDiff.gradient(f, [10.0, 10.0])
@show ForwardDiff.hessian(f, [10.0, 10.0])



Pkg.add("Ipopt")
if old
Pkg.add("JuMP")


function compute_friction(state, t, Δt, τ_propensity)
    m = Model(solver = IpoptSolver(print_level=0))
    JuMP.register(m, :mdp_friction, num_velocities(mechanism),
        (τ...) -> begin evaluate_friction_kinetic_energy(state, t, Δt, [τ...] + τ_propensity) end,
        autodiff=true)

    @variable(m, -10.0 <= τ[1:num_velocities(mechanism)] <= 10.0)
    @NLobjective(m, Min, mdp_friction(τ[1], τ[2]))

    solve(m)
    return getvalue(τ)
end

# can take ~40 seconds the first time
@show @time compute_friction(state, 0.0, 1e-3, [0.0, 0.0])

# output on ForwardDiff v0.5.0+
"""
  configuration(state) = [20.0, 0.0]
velocity(state) = [0.0, 0.0]
evaluate_friction_kinetic_energy(state, 0.0, 0.001, τ) = Dual{:some_tag}(0.04226111544136931,0.0005883280258112964)
ForwardDiff.gradient(f, [10.0, 10.0]) = [-0.00038079, 0.000969118]
ERROR: LoadError: MethodError: Cannot `convert` an object of type RigidBodyDynamics.CustomCollections.UnsafeFastDict{RigidBodyDynamics.Graphs.vertex_index,RigidBodyDynamics.RigidBody{Float64},Array{Array{RigidBodyDynamics.Contact.SoftContactStateDeriv{Void,RigidBodyDynamics.Contact.ViscoelasticCoulombStateDeriv{SubArray{ForwardDiff.Dual{ForwardDiff.Tag{##2#3,0x096989b645fcecd8},ForwardDiff.Dual{ForwardDiff.Tag{Void,0xb046287d533b082d},Float64,2},2},1,Array{ForwardDiff.Dual{ForwardDiff.Tag{##2#3,0x096989b645fcecd8},ForwardDiff.Dual{_,Float64,2},2},1},Tuple{UnitRange{Int64}},true}}},1},1} where _<:ForwardDiff.Tag} to an object of type RigidBodyDynamics.CustomCollections.UnsafeFastDict{RigidBodyDynamics.Graphs.vertex_index,RigidBodyDynamics.RigidBody{Float64},Array{Array{RigidBodyDynamics.Contact.SoftContactStateDeriv{Void,RigidBodyDynamics.Contact.ViscoelasticCoulombStateDeriv{SubArray{ForwardDiff.Dual{ForwardDiff.Tag{##2#3,0x096989b645fcecd8},ForwardDiff.Dual{ForwardDiff.Tag{Void,0xb046287d533b082d},Float64,2},2},1,Array{ForwardDiff.Dual{ForwardDiff.Tag{##2#3,0x096989b645fcecd8},ForwardDiff.Dual{ForwardDiff.Tag{Void,0xb046287d533b082d},Float64,2},2},1},Tuple{UnitRange{Int64}},true}}},1},1}}
This may have arisen from a call to the constructor RigidBodyDynamics.CustomCollections.UnsafeFastDict{RigidBodyDynamics.Graphs.vertex_index,RigidBodyDynamics.RigidBody{Float64},Array{Array{RigidBodyDynamics.Contact.SoftContactStateDeriv{Void,RigidBodyDynamics.Contact.ViscoelasticCoulombStateDeriv{SubArray{ForwardDiff.Dual{ForwardDiff.Tag{##2#3,0x096989b645fcecd8},ForwardDiff.Dual{ForwardDiff.Tag{Void,0xb046287d533b082d},Float64,2},2},1,Array{ForwardDiff.Dual{ForwardDiff.Tag{##2#3,0x096989b645fcecd8},ForwardDiff.Dual{ForwardDiff.Tag{Void,0xb046287d533b082d},Float64,2},2},1},Tuple{UnitRange{Int64}},true}}},1},1}}(...),
since type constructors fall back to convert methods.
Stacktrace:
 [1] RigidBodyDynamics.DynamicsResult{ForwardDiff.Dual{ForwardDiff.Tag{##2#3,0x096989b645fcecd8},ForwardDiff.Dual{ForwardDiff.Tag{Void,0xb046287d533b082d},Float64,2},2},M} where M(::RigidBodyDynamics.Mechanism{Float64}) at /Users/goretkin/playgrounds/jump/packages/v0.6/RigidBodyDynamics/src/dynamics_result.jl:73
 [2] evaluate_friction_kinetic_energy(::RigidBodyDynamics.MechanismState{Float64,Float64,Float64,TypeSortedCollections.TypeSortedCollection{Tuple{Array{RigidBodyDynamics.Joint{Float64,RigidBodyDynamics.Revolute{Float64}},1}},1}}, ::Float64, ::Float64, ::Array{ForwardDiff.Dual{ForwardDiff.Tag{##2#3,0x096989b645fcecd8},ForwardDiff.Dual{ForwardDiff.Tag{Void,0xb046287d533b082d},Float64,2},2},1}) at /Users/goretkin/projects/julia_lqr/min_jump.jl:63
 [3] vector_mode_gradient at /Users/goretkin/playgrounds/jump/packages/v0.6/ForwardDiff/src/gradient.jl:94 [inlined]
 [4] gradient(::##2#3, ::Array{ForwardDiff.Dual{ForwardDiff.Tag{Void,0xb046287d533b082d},Float64,2},1}, ::ForwardDiff.GradientConfig{ForwardDiff.Tag{##2#3,0x096989b645fcecd8},ForwardDiff.Dual{ForwardDiff.Tag{Void,0xb046287d533b082d},Float64,2},2,Array{ForwardDiff.Dual{ForwardDiff.Tag{##2#3,0x096989b645fcecd8},ForwardDiff.Dual{ForwardDiff.Tag{Void,0xb046287d533b082d},Float64,2},2},1}}) at /Users/goretkin/playgrounds/jump/packages/v0.6/ForwardDiff/src/gradient.jl:19
 [5] vector_mode_jacobian(::ForwardDiff.##43#44{##2#3,ForwardDiff.HessianConfig{ForwardDiff.Tag{##2#3,0x096989b645fcecd8},Float64,2,Array{ForwardDiff.Dual{ForwardDiff.Tag{##2#3,0x096989b645fcecd8},ForwardDiff.Dual{ForwardDiff.Tag{Void,0xb046287d533b082d},Float64,2},2},1},0xb046287d533b082d,Array{ForwardDiff.Dual{ForwardDiff.Tag{Void,0xb046287d533b082d},Float64,2},1}}}, ::Array{Float64,1}, ::ForwardDiff.JacobianConfig{ForwardDiff.Tag{Void,0xb046287d533b082d},Float64,2,Array{ForwardDiff.Dual{ForwardDiff.Tag{Void,0xb046287d533b082d},Float64,2},1}}) at /Users/goretkin/playgrounds/jump/packages/v0.6/ForwardDiff/src/jacobian.jl:132
 [6] jacobian(::ForwardDiff.##43#44{##2#3,ForwardDiff.HessianConfig{ForwardDiff.Tag{##2#3,0x096989b645fcecd8},Float64,2,Array{ForwardDiff.Dual{ForwardDiff.Tag{##2#3,0x096989b645fcecd8},ForwardDiff.Dual{ForwardDiff.Tag{Void,0xb046287d533b082d},Float64,2},2},1},0xb046287d533b082d,Array{ForwardDiff.Dual{ForwardDiff.Tag{Void,0xb046287d533b082d},Float64,2},1}}}, ::Array{Float64,1}, ::ForwardDiff.JacobianConfig{ForwardDiff.Tag{Void,0xb046287d533b082d},Float64,2,Array{ForwardDiff.Dual{ForwardDiff.Tag{Void,0xb046287d533b082d},Float64,2},1}}) at /Users/goretkin/playgrounds/jump/packages/v0.6/ForwardDiff/src/jacobian.jl:21
 [7] hessian(::##2#3, ::Array{Float64,1}) at /Users/goretkin/playgrounds/jump/packages/v0.6/ForwardDiff/src/hessian.jl:19
 [8] include_from_node1(::String) at ./loading.jl:569
 [9] include(::String) at ./sysimg.jl:14
"""

tab_it(s) = "\t" * replace(s, "\n", "\n\t")

# like the version in Base, but keeps the splitters
function _split(str::AbstractString, splitter, strs::Array)
    i = start(str)
    n = endof(str)
    r = search(str,splitter,i)
    j, k = first(r), nextind(str,last(r))
    while 0 < j <= n
        if i < k
            if i <= prevind(str,j) # don't add empty strings *shrug*
              push!(strs, SubString(str,i,prevind(str,j)))
            end
            push!(strs, SubString(str,first(r),last(r)))
            i = k
        end
        (k <= j) && (k = nextind(str,j))
        r = search(str,splitter,k)
        j, k = first(r), nextind(str,last(r))
    end
    push!(strs, SubString(str,i))
    return strs
end

function format_type(s)
  splitted = _split(s, ['{', '}', ','], [])
  nesting = cumsum(map(
              s-> if s=="{" return 1; elseif s=="}" return -1; else return 0; end,
              splitted))
  out = ""

  for (n, s) in zip(nesting, splitted)
    prefix = (s in ["{", "}", ","]) ? "" : ("\n" * repeat("  ", n))
    out *= (prefix * s)
  end
  return out
end
  
# prettier types
"""
from = format_type( "RigidBodyDynamics.CustomCollections.UnsafeFastDict{RigidBodyDynamics.Graphs.vertex_index,RigidBodyDynamics.RigidBody{Float64},Array{Array{RigidBodyDynamics.Contact.SoftContactStateDeriv{Void,RigidBodyDynamics.Contact.ViscoelasticCoulombStateDeriv{SubArray{ForwardDiff.Dual{ForwardDiff.Tag{##2#3,0x096989b645fcecd8},ForwardDiff.Dual{ForwardDiff.Tag{Void,0xb046287d533b082d},Float64,2},2},1,Array{ForwardDiff.Dual{ForwardDiff.Tag{##2#3,0x096989b645fcecd8},ForwardDiff.Dual{_,Float64,2},2},1},Tuple{UnitRange{Int64}},true}}},1},1} where _<:ForwardDiff.Tag}")
  
to = format_type( "RigidBodyDynamics.CustomCollections.UnsafeFastDict{RigidBodyDynamics.Graphs.vertex_index,RigidBodyDynamics.RigidBody{Float64},Array{Array{RigidBodyDynamics.Contact.SoftContactStateDeriv{Void,RigidBodyDynamics.Contact.ViscoelasticCoulombStateDeriv{SubArray{ForwardDiff.Dual{ForwardDiff.Tag{##2#3,0x096989b645fcecd8},ForwardDiff.Dual{ForwardDiff.Tag{Void,0xb046287d533b082d},Float64,2},2},1,Array{ForwardDiff.Dual{ForwardDiff.Tag{##2#3,0x096989b645fcecd8},ForwardDiff.Dual{ForwardDiff.Tag{Void,0xb046287d533b082d},Float64,2},2},1},Tuple{UnitRange{Int64}},true}}},1},1}}")

print(from)

RigidBodyDynamics.CustomCollections.UnsafeFastDict{
  RigidBodyDynamics.Graphs.vertex_index,
  RigidBodyDynamics.RigidBody{
    Float64},
  Array{
    Array{
      RigidBodyDynamics.Contact.SoftContactStateDeriv{
        Void,
        RigidBodyDynamics.Contact.ViscoelasticCoulombStateDeriv{
          SubArray{
            ForwardDiff.Dual{
              ForwardDiff.Tag{
                ##2#3,
                0x096989b645fcecd8},
              ForwardDiff.Dual{
                ForwardDiff.Tag{
                  Void,
                  0xb046287d533b082d},
                Float64,
                2},
              2},
            1,
            Array{
              ForwardDiff.Dual{
                ForwardDiff.Tag{
                  ##2#3,
                  0x096989b645fcecd8},
                ForwardDiff.Dual{
                  _,          # this line differs
                  Float64,
                  2},
                2},
              1},
            Tuple{
              UnitRange{
                Int64}},
            true}}},
      1},
    1}
   where _<:ForwardDiff.Tag}  # this line differs

print(to)


RigidBodyDynamics.CustomCollections.UnsafeFastDict{
  RigidBodyDynamics.Graphs.vertex_index,
  RigidBodyDynamics.RigidBody{
    Float64},
  Array{
    Array{
      RigidBodyDynamics.Contact.SoftContactStateDeriv{
        Void,
        RigidBodyDynamics.Contact.ViscoelasticCoulombStateDeriv{
          SubArray{
            ForwardDiff.Dual{
              ForwardDiff.Tag{
                ##2#3,
                0x096989b645fcecd8},
              ForwardDiff.Dual{
                ForwardDiff.Tag{
                  Void,
                  0xb046287d533b082d},
                Float64,
                2},
              2},
            1,
            Array{
              ForwardDiff.Dual{
                ForwardDiff.Tag{
                  ##2#3,
                  0x096989b645fcecd8},
                ForwardDiff.Dual{
                  ForwardDiff.Tag{      # there is a concrete tag
                    Void,
                    0xb046287d533b082d},
                  Float64,
                  2},
                2},
              1},
            Tuple{
              UnitRange{
                Int64}},
            true}}},
      1},
    1}}
"""