[WIP] Add ability to constrain the orientation of mechanism links

Author: ferrolho

Project.toml

@@ -18,6 +18,7 @@ PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Requires = "ae029012-a4dd-5104-9daa-d747884805df"
 RigidBodyDynamics = "366cf18f-59d5-5db9-a4de-86a9f6786172"
+Rotations = "6038ab10-8711-5258-84ad-4b1120ba62dc"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 SparseDiffTools = "47a9eef4-7e08-11e9-0b38-333d64bd3804"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
@@ -34,6 +35,7 @@ Plots = "1"
 PrecompileTools = "1"
 Requires = "1"
 RigidBodyDynamics = "2"
+Rotations = "1"
 SparseDiffTools = "2"
 StaticArrays = "1"
 julia = "1.10"

interoperability/python_calls_julia.py

@@ -12,6 +12,7 @@ def init():
     """
     # Load package dependencies
     jl.seval("using MeshCat")
+    jl.seval("using Rotations")
     jl.seval("using TORA")
 
     # Call the greet() function from the TORA.jl package
@@ -55,7 +56,7 @@ def main():
 
     jl_robot = init()
 
-    trajectory_discretisation = 100  # in Hertz
+    trajectory_discretisation = 10  # in Hertz
     trajectory_duration = 2.0  # in seconds
 
     trajectory_num_knots = int(trajectory_discretisation * trajectory_duration) + 1  # in units (number of knots)
@@ -95,11 +96,18 @@ def main():
         jl.TORA.fix_joint_velocities_b(jl_problem, jl_robot, jl_problem.num_knots, np.zeros(jl_robot.n_v))
 
         # Define the end-effector position that we want to achieve
-        py_eff_pos = [0.5, 0.2, 0.1]
+        py_eff_pos = [0.5, 0.2, 0.8]
 
         # Constrain the end-effector position at the last knot of the trajectory
         jl.TORA.constrain_ee_position_b(jl_problem, jl_problem.num_knots, py_eff_pos)
 
+        # Define the end-effector orientation that we want to achieve
+        body_name = "panda_link7"  # this is the fixed parent link of "panda_hand_tcp"
+        jl_quaternion = jl.QuatRotation(0.27, 0.65, 0.27, 0.65)  # Quaternion order is (w, x, y, z)
+
+        # Constrain the end-effector orientation at the last knot of the trajectory
+        jl.TORA.add_constraint_body_orientation_b(jl_problem, jl_robot, body_name, jl_problem.num_knots, jl_quaternion)
+
         # Show a summary of the problem instance (this can be commented out)
         jl.TORA.show_problem_info(jl_problem)
 

interoperability/python_setup.py

@@ -11,6 +11,7 @@
 
 # Add package dependencies
 jl.seval('Pkg.add("MeshCat")')
+jl.seval('Pkg.add("Rotations")')
 
 # Develop the local package of TORA.jl
 jl.seval('Pkg.develop(path="../../TORA.jl")')

src/TORA.jl

@@ -10,6 +10,7 @@ using NPZ
 using Plots
 using Requires
 using RigidBodyDynamics
+using Rotations
 using SparseArrays
 using SparseDiffTools
 using StaticArrays

src/constraints/end_effector.jl

@@ -64,3 +64,17 @@ function show_ee_path(vis::Visualizer, ee_positions::Matrix{Float64})
     setobject!(vis["ee path"], Object(PointCloud(points), material, "Line"))
     nothing
 end
+
+function body_orientation!(dx, robot, x::AbstractVector{T}) where {T}
+    state = robot.statecache[T]
+    set_configuration!(state, x)
+
+    body_tf = transform_to_root(state, robot.frame_ee)
+    mrp = MRP(rotation(body_tf))
+
+    dx[1] = mrp.x
+    dx[2] = mrp.y
+    dx[3] = mrp.z
+
+    dx
+end

src/precompile.jl

@@ -39,6 +39,10 @@ using PrecompileTools: @setup_workload, @compile_workload
         fix_joint_velocities!(problem, robot, 1, zeros(robot.n_v))
         constrain_ee_position!(problem, 1, zeros(3))
 
+        body_name = "panda_link7"  # this is the fixed parent link of "panda_hand_tcp"
+        quaternion = QuatRotation(0.27, 0.65, 0.27, 0.65)
+        add_constraint_body_orientation!(problem, robot, body_name, 1, quaternion)
+
         initial_guess = zeros(problem.num_knots * (robot.n_q + robot.n_v + robot.n_τ) - robot.n_τ)
 
         cpu_time, x, solver_log = solve_with_ipopt(problem, robot, initial_guess=initial_guess, user_options=ipopt_options, use_inv_dyn=false, minimise_τ=false)

src/problem.jl

@@ -22,6 +22,9 @@ mutable struct Problem
 
     ee_pos::Dict{Int64,Point{3,Float64}}  # End-effector target positions
 
+    # Dictionaries indexing: knot → body_name → value
+    constraints_body_orientation::Dict{Int64,Dict{String,Rotation}}
+
     # Jacobian data (e.g., for sparse Jacobian of dynamics with AD)
     const jacdata_fwd_dyn::JacobianData
     const jacdata_inv_dyn::JacobianData
@@ -58,6 +61,7 @@ mutable struct Problem
             Dict{Int64,Vector{Float64}}(),
             Dict{Int64,Vector{Float64}}(),
             Dict{Int64,Point{3,Float64}}(),
+            Dict{Int64,Dict{String,Rotation}}(),
             jacdata_fwd_dyn,
             jacdata_inv_dyn,
             jacdata_ee_position
@@ -121,6 +125,55 @@ function constrain_ee_position!(problem::Problem, knot, position)
     return problem
 end
 
+"""
+    add_constraint_body_orientation!(problem, robot, body_name, knot, rotation; force=false)
+
+Constrain the orientation of a body (kinematic link) of the robot.
+
+See also: [`add_constraint_body_position!`](@ref)
+"""
+function add_constraint_body_orientation!(problem::Problem, robot::Robot, body_name::String, knot::Integer, rotation::Rotation; force=false)
+    @assert body_name ∈ map(x -> x.name, bodies(robot.mechanism))
+    @assert 1 <= knot <= problem.num_knots
+
+    d = problem.constraints_body_orientation
+
+    if !haskey(d, knot)
+        d[knot] = Dict{String,Rotation}()
+    end
+
+    dₖ = d[knot]
+
+    if !haskey(dₖ, body_name) || force
+        dₖ[body_name] = rotation
+    elseif haskey(dₖ, body_name)
+        @warn """
+        You tried to add a constraint for `$(body_name)` at knot=$(knot),
+        but you have already defined a constraint for that knot before.
+        If you wish to overwrite previous constraints, use `force=true`.
+        """
+    end
+
+    return problem
+end
+
+"""
+    body_orientation_constraints_per_knot(problem)
+
+Count number of transcription constraints needed to enforce the
+orientation constraints of the high-level problem definition.
+"""
+function body_orientation_constraints_per_knot(problem::Problem)
+    map(1:problem.num_knots) do knot
+        counter = 0
+        dₖ = get(problem.constraints_body_orientation, knot, Dict())
+        for (body_name, rotation) ∈ dₖ
+            counter += 3  # MRP values
+        end
+        counter
+    end
+end
+
 """
     show_problem_info(problem)
 
@@ -129,12 +182,15 @@ Output a summary of the problem, including the number of knots with constraints.
 function show_problem_info(problem::Problem)
     t = (problem.num_knots - 1) * problem.dt
 
+    println("Discretization ....................................... $(problem.num_knots) knots")
     println("Motion duration ...................................... $(t) seconds")
-    println("Number of knots ...................................... $(problem.num_knots)")
+    println("Time step length ..................................... $(problem.dt) seconds")
+    println()
     println("Number of knots with constrained joint positions ..... $(length(problem.fixed_q))")
     println("Number of knots with constrained joint velocities .... $(length(problem.fixed_v))")
     println("Number of knots with constrained joint torques ....... $(length(problem.fixed_τ))")
     println("Number of knots with constrained ee position ......... $(length(problem.ee_pos))")
+    println("Number of knots with body orientation constraints .... $(length(problem.constraints_body_orientation))")
 end
 
 """

src/transcription/ipopt.jl

@@ -73,11 +73,13 @@ function solve_with_ipopt(problem::Problem, robot::Robot;
 
     use_m₁ = true  # nonlinear equality: dynamics
     use_m₂ = true  # nonlinear equality: ee xyz-position
+    use_m₃ = true  # nonlinear equality: ee rotation
 
     m₁ = !use_m₁ ? 0 : (robot.n_q + robot.n_v) * (problem.num_knots - 1)  # equations of motion
     m₂ = !use_m₂ ? 0 : 3 * length(problem.ee_pos)  # xyz-position (3 NL equality)
+    m₃ = !use_m₃ ? 0 : body_orientation_constraints_per_knot(problem) |> sum
 
-    m = m₁ + m₂  # total number of constraints
+    m = m₁ + m₂ + m₃  # total number of constraints
 
     g_L, g_U = Float64[], Float64[]
 
@@ -93,7 +95,28 @@ function solve_with_ipopt(problem::Problem, robot::Robot;
         append!(g_U, vec(con_ee))
     end
 
-    @assert length(g_L) == length(g_U)
+    body_name = "panda_link7"  # TODO ... hard-coded for now.
+    # We want to iterate only through the constrained knots for this robot body/link.
+    knots = filter(1:problem.num_knots) do knotᵢ
+        knotᵢ ∈ keys(problem.constraints_body_orientation) &&
+        body_name ∈ keys(problem.constraints_body_orientation[knotᵢ])
+    end
+
+    jacdata = JacobianData((out, x) -> body_orientation!(out, robot, x), rand(3), rand(robot.n_q))
+
+    if use_m₃
+        # These are the target values for the equalities of Knitro.
+        con_vals = mapreduce(hcat, knots; init=Matrix{Float64}(undef, 3, 0)) do k
+            dₖ = problem.constraints_body_orientation[k]
+            mrp = MRP(dₖ[body_name])
+            [mrp.x, mrp.y, mrp.z]
+        end
+
+        append!(g_L, con_vals)
+        append!(g_U, con_vals)
+    end
+
+    @assert length(g_L) == length(g_U) == m
     @assert eltype(g_L) == eltype(g_U) == Float64
 
     # dimension of each mesh point
@@ -122,14 +145,25 @@ function solve_with_ipopt(problem::Problem, robot::Robot;
         end
 
         if use_m₂
-            length_c = 3
+            length_c = size(problem.jacdata_ee_position.jac, 1)  # length of constraint (per knot)
             for (i, k) = enumerate(sort(collect(keys(problem.ee_pos))))
                 ind_vars = range(1 + (k - 1) * nₓ, length=robot.n_q)  # indices of the decision variables
                 offset_con = (i - 1) * length_c + (m₁)
                 ind_cons = (1:length_c) .+ offset_con
                 @views ee_position!(g[ind_cons], robot, x[ind_vars])  # constraint evaluation at that point
             end
         end
+
+        if use_m₃
+            length_c = size(jacdata.jac, 1)  # length of constraint (per knot)
+            offset_con = m₁ + m₂
+            for i = knots .- 1
+                ind_cons = (1:length_c) .+ offset_con                      # Indices of the respective constraints
+                ind_vars = range(1 + i * nₓ, length=robot.n_q)             # Indices of the decision variables
+                @views body_orientation!(g[ind_cons], robot, x[ind_vars])  # Constraint evaluation at that point
+                offset_con += length_c
+            end
+        end
     end
 
     jacIndexConsCB = Cint[]
@@ -167,6 +201,22 @@ function solve_with_ipopt(problem::Problem, robot::Robot;
         @assert length(jacIndexConsCB) == length(jacIndexVarsCB)
     end
 
+    if use_m₃
+        ind_offset = m₁ + m₂
+        for k = knots
+            inds = rowvals(jacdata.jac) .+ ind_offset
+            ind_offset += size(jacdata.jac, 1)
+            append!(jacIndexConsCB, inds)
+        end
+
+        for k = knots
+            vals = vcat([fill(j + (k - 1) * nₓ, length(nzrange(jacdata.jac, j)))
+                         for j = 1:size(jacdata.jac, 2)]...)
+            append!(jacIndexVarsCB, vals)
+        end
+
+        @assert length(jacIndexConsCB) == length(jacIndexVarsCB)
+    end
 
     function eval_jac_g(x, rows, cols, values::Union{Nothing,Vector{Float64}})
         if isnothing(values)
@@ -186,22 +236,36 @@ function solve_with_ipopt(problem::Problem, robot::Robot;
                     # Evaluate the Jacobian at that point
                     jacdata_dyn(x[ind_vars])
 
-                    # Pass the Jacobian to Knitro
-                    offset_jac = i * jacdata_dyn.jac_length
-                    ind_jac = (1:jacdata_dyn.jac_length) .+ offset_jac
+                    # Pass the Jacobian to Ipopt
+                    ind_jac = (1:jacdata_dyn.jac_length) .+ offset_prev
                     values[ind_jac] = nonzeros(jacdata_dyn.jac)
+                    offset_prev += jacdata_dyn.jac_length
                 end
-                offset_prev += (problem.num_knots - 1) * jacdata_dyn.jac_length
             end
 
             if use_m₂
                 # End-effector constraints
                 for (i, k) = enumerate(sort(collect(keys(problem.ee_pos))))
                     ind_qᵢ = range(1 + (k - 1) * nₓ, length=robot.n_q)  # indices of the decision variables
                     problem.jacdata_ee_position(x[ind_qᵢ])              # jacobian evaluation at that point
-                    offset_jac = (i - 1) * problem.jacdata_ee_position.jac_length + offset_prev
-                    ind_jac = (1:problem.jacdata_ee_position.jac_length) .+ offset_jac
-                    values[ind_jac] = nonzeros(problem.jacdata_ee_position.jac)  # pass jacobian to Knitro
+
+                    # Pass the Jacobian to Ipopt
+                    ind_jac = (1:problem.jacdata_ee_position.jac_length) .+ offset_prev
+                    values[ind_jac] = nonzeros(problem.jacdata_ee_position.jac)
+                    offset_prev += problem.jacdata_ee_position.jac_length
+                end
+            end
+
+            if use_m₃
+                # End-effector constraints (rotation)
+                for k = knots
+                    ind_qᵢ = range(1 + (k - 1) * nₓ, length=robot.n_q)  # Indices of the decision variables
+                    jacdata(x[ind_qᵢ])  # Evaluate the Jacobian at that point
+
+                    # Pass the Jacobian to Ipopt
+                    ind_jac = (1:jacdata.jac_length) .+ offset_prev
+                    values[ind_jac] = nonzeros(jacdata.jac)
+                    offset_prev += jacdata.jac_length
                 end
             end
         end
@@ -244,6 +308,7 @@ function solve_with_ipopt(problem::Problem, robot::Robot;
     nele_jac = 0
     nele_jac += !use_m₁ ? 0 : jacdata_dyn.jac_length * (problem.num_knots - 1)
     nele_jac += !use_m₂ ? 0 : problem.jacdata_ee_position.jac_length * length(problem.ee_pos)
+    nele_jac += !use_m₃ ? 0 : jacdata.jac_length * length(knots)
 
     prob = Ipopt.CreateIpoptProblem(n, vec(x_L), vec(x_U), m, g_L, g_U, nele_jac, 0,
                                     eval_f, eval_g, eval_grad_f, eval_jac_g, nothing)