Humanoid LQR

Our second example is a translation of DeepMind's LQR tutorial into Julia. The aim is to balance a humanoid on one leg with a Linear Quadratic Regulator (LQR). Users familiar with MuJoCo's python bindings might find it useful to compare this example to the DeepMind tutorial to see the similarities and differences with MuJoCo.jl.

The humanoid model

Let's start by having a look at the humanoid model shipped with MuJoCo. You can find a copy of the humanoid.xml file here, or locally in the directory given by running example_model_files_directory.

using MuJoCo

model = load_model(joinpath(example_model_files_directory(), "humanoid", "humanoid.xml"))
data = init_data(model)

The model comes shipped with three different "keyframes" defining particular starting poses of the humanoid. You can inspect each one by re-setting the model to a specific keyframe with resetkey! and running visualise!.

init_visualiser()

for i in 1:3
    resetkey!(model, data, i)
    visualise!(model, data)
end

None of these initial states are stable. In this example, we'll focus on designing a controller for the second keyframe to get the humanoid to stand and balance on one leg.

Computing the control set-point

To design our linear controller, we'll need a set-point $(x^*, u^*)$ about which to stabilise the system. We already have $x^*$, since the desired state is just the starting position of the humanoid in keyframe 2. To find the control inputs required to hold the humanoid in this position, we'll use inverse dynamics via MuJoCo's mj_inverse. If we set all joint accelerations to zero, mj_inverse computes the forces required to hold the model in a given configuration. Let's try it out.

# Reset to desired keyframe
keyframe = 2
resetkey!(model, data, keyframe)

# Propagate derived quantities
mj_forward(model, data)
# Set joint accelerations to 0
data.qacc .= 0

# Inspect forces from inverse dynamics
mj_inverse(model, data)
println("Required control: ", data.qfrc_inverse)
Required control: [0.0; 0.0; 275.87856563199057; -33.18553410706376; 4.995479129525545; -6.6877127833417545; -4.305208663659141; 3.6932410280923613; -15.451150947683784; -10.906299621526648; 0.4118355495043427; -1.6130054371234719; -9.792551439832716; -2.311829304174266; -0.36610644238969153; -5.9125502445456775; -0.41711259718163185; -1.914082279494563; 5.758907794560738; 2.665187588315114; -0.2020663154044419; -5.754549922711961; 0.9943292302894648; 1.141092437180831; -1.9874024238883508; 3.821127395975936; 1.1510540373906009;;]

Unfortunately for us, things are not so simple. data.qfrc_inverse stores the calculated forces that are required to achieve the desired joint acceleration. Looking at its output, we see that the third element is a large value. This corresponds to an unphysical vertical force due to the humanoid starting slightly above the ground. The simulator has been forced to introduce this impossible force to meet the requirement that data.qacc == 0.

We can remove this unphysical force by slightly tweaking the starting height of the model. Let's examine how this large force changes when we move the humanoid up/down from its starting position. We will examine the range of heights from $\pm 1$mm. Note that all heights are all measured in metres.

using CairoMakie
using LinearAlgebra

heights = LinRange(-0.001, 0.001, 2001) # -1mm to +1mm
# Map each height to the corresponding high force output
u_vert = map(heights) do h
    # Set model in position with qacc == 0
    resetkey!(model, data, keyframe)
    mj_forward(model, data)
    data.qacc .= 0

    # Offset the height and check required vertical forces
    data.qpos[3] += h
    mj_inverse(model, data)
    return data.qfrc_inverse[3] # 3 -> z-force
end

# Find height corresponding to minimum fictitious force (best offset)
height = heights[argmin(abs.(u_vert))]
height_mm = height*1000
heights_mm = heights .* 1000

# Compare force to weight of humanoid
weight = sum(model.body_mass) * norm(model.opt.gravity)

fig = Figure(resolution=(500,300))
ax = Axis(fig[1,1], xlabel="Vertical offset (mm)", ylabel="Vertical force (N)")
lines!(ax, heights_mm, u_vert, label="Vertical force")
lines!(ax, heights_mm, weight*ones(length(heights)), linestyle=:dash, label="Weight")
lines!(ax, [height_mm, height_mm], [minimum(u_vert), maximum(u_vert)], linestyle=:dash)
axislegend(ax,position=:rb)
┌ Warning: Found `resolution` in the theme when creating a `Scene`. The `resolution` keyword for `Scene`s and `Figure`s has been deprecated. Use `Figure(; size = ...` or `Scene(; size = ...)` instead, which better reflects that this is a unitless size and not a pixel resolution. The key could also come from `set_theme!` calls or related theming functions.
└ @ Makie ~/.julia/packages/Makie/rEu75/src/scenes.jl:227

As we lift the humanoid upwards, the vertical force increases until it exactly balances the weight of the humanoid. This is the force required to keep the accelerations zero while the model is in the air. Pushing the humanoid down, however, requires an increasingly negative vertical force to explain the fact that its foot is penetrating the floor but data.qacc == 0. At a height of approximately $-0.5\,$mm the humanoid's foot is just resting on the floor and all forces can be accounted for by internal torques only. This is the case we're interested in.

Let's now take our best-choice offset and save the humanoid's position and forces (from inverse dynamics) to define our set point. Our position set-point is just data.qpos after changing the height.

resetkey!(model, data, keyframe)
mj_forward(model, data)
data.qacc .= 0
data.qpos[3] += height
qpos0 = copy(data.qpos)
println(qpos0)
[0.0; 0.0; 1.2189729999999999; 0.971587999157264; -0.17997299984389503; 0.13531799988262788; -0.07290759993676141; -0.0516; -0.202; 0.23; -0.24; -0.007; -0.34; -1.76; -0.466; -0.0415; -0.08; -0.01; -0.37; -0.685; -0.35; -0.09; 0.109; -0.067; -0.7; -0.05; 0.12; 0.16;;]

We can compute the control set-point by propagating the required forces through the actuator dynamics. For the simple linear motors in our humanoid robot, we can directly use data.actuator_moment as the Jacobian of data.qfrc_inverse with respect to data.ctrl.

mj_inverse(model, data)
qfrc0 = copy(data.qfrc_inverse)

M_act = data.actuator_moment
ctrl0 = pinv(M_act)' * qfrc0
println(ctrl0)
[0.06464056026858045; 0.36593038170982806; -0.4660050472258556; -0.2726574905381662; 0.010295888737608568; -0.013441711976028932; -0.12240689299790895; -0.11559146520871331; -0.018305322119484577; -0.5938650629988403; -0.0542756185108938; 0.10219803776978083; 0.326260284296838; 0.6668428405823527; -0.005652473252838193; -0.28772749613559806; 0.049716461514473245; 0.02852731092952078; -0.09937012119441754; 0.1910563697987968; 0.02877635093476502;;]

As a sanity check, let's verify that these control inputs actually give us the right forces.

data.ctrl .= ctrl0
mj_forward(model, data)
qfrc_test = data.qfrc_actuator
println("Desired forces: ", qfrc0)
println("Actual forces:  ", qfrc_test)
println("Joint forces equal? ", all((qfrc_test .≈ qfrc0)[7:end]))
Desired forces: [0.0; 0.0; -0.19089622146884722; -3.447398548108275; 0.22200083297670048; -0.8168220883582915; 2.585622410743218; 14.63721526839312; -18.640201889034223; -10.906299621526648; 0.4118355495043427; -1.6130054371234719; -9.792551439832716; -2.311829304174266; -0.36610644238969153; -23.75460251995361; -2.171024740435752; 12.263764532373699; 26.100822743747038; 13.336856811647053; -0.11304946505676386; -5.754549922711961; 0.9943292302894648; 1.141092437180831; -1.9874024238883508; 3.821127395975936; 1.1510540373906009;;]
Actual forces:  [0.0; 0.0; 0.0; 0.0; 0.0; 0.0; 2.585622410743218; 14.637215268393122; -18.640201889034223; -10.906299621526648; 0.4118355495043427; -1.6130054371234719; -9.792551439832716; -2.311829304174266; -0.36610644238969153; -23.754602519953615; -2.171024740435752; 12.263764532373699; 26.10082274374704; 13.336856811647053; -0.11304946505676386; -5.754549922711961; 0.9943292302894648; 1.141092437180831; -1.987402423888351; 3.821127395975936; 1.1510540373906009;;]
Joint forces equal? true

The actuator forces are approximately equal to the desired forces qfrc0 for all joints in the humanoid. The first six elements correspond to forces on the "root joint" (the free body of the humanoid). There is still a slight mismatch here, but it is minor. Let's also have a look at how the model behaves if we start it exactly at our set point.

reset!(model, data)
data.qpos .= qpos0
data.ctrl .= ctrl0
visualise!(model, data)

The humanoid still falls over because we are trying to stabilise the system at an unstable equilibrium point, and the control set-point ctrl0 is only an approximation. Even if we had found the exact controls to hold the system in qpos0, any slight perturbation would cause the humanoid to fall over. We therefore need a controller.

Designing the LQR cost

See A quick review of LQR if you're in need of a refresher on linear quadratic regulators.

Now that we have our set point to stabilise, we'll need to design the LQR weight matrices $Q,R$ to encourage the system to remain balanced. Let's start by setting R equal to the identity matrix I and defining some useful variables.

nu = model.nu # number of actuators/controls
nv = model.nv # number of degrees of freedom
R = Matrix{Float64}(I, nu, nu)

We'll construct the $Q$ matrix in two parts: one component to encourage the humanoid's centre of mass (CoM) to remain above its foot (helps with stability), and another component to keep the joints close to their original configuration qpos0.

CoM balancing cost

Much like the Python bindings, MuJoCo.jl includes a number of useful functions to isolate and work with different parts of a MuJoCo model. Let's extract references for the torso and left foot and have a look at one of them.

import MuJoCo as MJ

torso = MJ.body(model, "torso")
left_foot = MJ.body(model, "foot_left")
ModelBody:
id:           10
name:         "foot_left"
parentid:     (1 x  )	Int32	[9]
rootid:       (1 x  )	Int32	[1]
weldid:       (1 x  )	Int32	[10]
mocapid:      (1 x  )	Int32	[-1]
jntnum:       (1 x  )	Int32	[2]
jntadr:       (1 x  )	Int32	[14]
dofnum:       (1 x  )	Int32	[2]
dofadr:       (1 x  )	Int32	[19]
geomnum:      (1 x  )	Int32	[2]
geomadr:      (1 x  )	Int32	[12]
simple:       (1 x  )	UInt8	[0]
sameframe:    (1 x  )	UInt8	[0]
pos:          (3 x  )	Float64	[0.0, 0.0, -0.39] (4 s.f.)
quat:         (4 x  )	Float64	[1.0, 0.0, 0.0, 0.0] (4 s.f.)
ipos:         (3 x  )	Float64	[0.035, 0.0, 0.0] (4 s.f.)
iquat:        (4 x  )	Float64	[0.0, 0.7071, 0.0, 0.7071] (4 s.f.)
mass:         (1 x  )	Float64	[1.131] (4 s.f.)
subtreemass:  (1 x  )	Float64	[1.131] (4 s.f.)
inertia:      (3 x  )	Float64	[0.006446, 0.005944, 0.000903] (4 s.f.)
invweight0:   (2 x  )	Float64	[0.3657, 48.34] (4 s.f.)
user:         nothing

One way to keep the CoM over the left foot is to design $Q$ from the difference between the CoM and left-foot Jacobians. MuJoCo's C API comes with a number of very useful functions for computing Jacobians.

# Get Jacobian for torso CoM
reset!(model, data)
data.qpos .= qpos0
forward!(model, data)
jac_com = mj_zeros(3, nv)
mj_jacSubtreeCom(model, data, jac_com, torso.id)

# Get (left) foot Jacobian for balancing
jac_foot = mj_zeros(3, nv)
mj_jacBodyCom(model, data, jac_foot, nothing, left_foot.id)

# Design Q-matrix to balance CoM over foot
jac_diff = jac_com .- jac_foot
Qbalance = jac_diff' * jac_diff
[4.930380657631324e-32 0.0 0.0 -3.4455790699220185e-17 1.6435007843370364e-16 6.068279684531846e-17 5.957832350481518e-17 1.4884641559545154e-16 -2.6158892480824602e-17 1.844872617411259e-18 -4.196752724496389e-19 -1.259491296360777e-17 6.743962466892438e-19 -4.265016092358265e-19 3.9565163223496944e-20 2.4416066810159544e-17 2.8506203702269416e-18 1.48013817644968e-16 -8.060974977015126e-17 1.7108492080947968e-17 1.2700948868758066e-18 -6.696119442090515e-19 3.3078500592035365e-18 1.3022159717762597e-19 -6.384933325393985e-19 1.5643617507679333e-18 -1.0971943724341637e-18; 0.0 4.930380657631324e-32 0.0 -1.7026868264106962e-16 -1.6855275442690524e-17 -5.1006340013404474e-17 -5.672925244624975e-17 -1.9676320315043004e-17 -1.5233218175986235e-16 1.1823574006427848e-17 -5.416746242122389e-19 2.476407050457178e-18 -2.0234689214836416e-18 1.8479027549073231e-19 2.8377258402809243e-19 1.467916473731728e-16 1.2536739852398554e-17 -2.2479389061965568e-17 1.389303693462475e-17 -2.0541067993702896e-18 1.1030517252671916e-17 1.960313318173584e-18 2.1048330148473913e-18 1.0926750149511406e-18 -2.2505531625316204e-18 -1.4222930571743299e-18 6.97275397205803e-21; 0.0 0.0 4.930380657631324e-32 -1.740532938265174e-18 -8.265421392186514e-19 -7.983078331039036e-19 1.658735815403105e-18 8.285949303260346e-18 -1.1446642093024695e-17 -4.7139329501991025e-18 2.670186302440841e-19 9.902981233405554e-19 -4.451396093616671e-18 2.6831269742696113e-19 -1.3389095381485578e-19 -1.2870466564366527e-17 -1.0696669620797387e-18 7.827037144299548e-18 1.5850258144832543e-17 7.536384440872769e-18 -4.655899830429835e-21 -3.2494013161563927e-18 5.881564313829157e-19 6.323587785736059e-19 -1.0736495926678668e-18 2.051052415391445e-18 6.378791949185355e-19; -3.4455790699220185e-17 -1.7026868264106962e-16 -1.740532938265174e-18 0.612156694417241 -0.056617286346724645 0.13376853659664892 0.1542174992158436 -0.0363619341993839 0.5447581101543878 -0.04195510048875402 0.002154513427850427 0.00021478001431383942 0.0066738119442117695 -0.00034957855947544386 -0.0010029203595738237 -0.5235476858711073 -0.04524950348585281 -0.026083592007825382 0.007795253550627904 -0.005128486145404674 -0.038980877842172516 -0.006187195034468947 -0.009601397385631415 -0.003886837103210349 0.008256304443286854 0.0037461756535838095 0.0007201717420216381; 1.6435007843370364e-16 -1.6855275442690524e-17 -8.265421392186514e-19 -0.056617286346724645 0.5536232120550901 0.21973167831675927 0.21796528600528056 0.5027547143173895 -0.034929433164675616 0.002186679578873358 -0.0012182485546882798 -0.04284727947400572 0.0030144217826222086 -0.001489378625558807 3.7119631416326994e-5 0.03142164746853996 0.005234354521605073 0.5009452854191183 -0.2737210674981011 0.057605601051604054 0.00046287776926983907 -0.002847784767902927 0.010297009006599165 4.993367362317162e-5 -0.0013409779850553514 0.005666516559787287 -0.003670482168746765; 6.068279684531846e-17 -5.1006340013404474e-17 -7.983078331039036e-19 0.13376853659664892 0.21973167831675927 0.12746857263854613 0.13198994379348716 0.20342078717952894 0.12558158636649938 -0.009884876441425145 3.9522459447686384e-5 -0.018079690582339526 0.002995459962873844 -0.0007204511116886534 -0.0002427072224056781 -0.12160105430170758 -0.009443807843236254 0.20530332455179046 -0.11384336848933328 0.023060030113865867 -0.009848117628120642 -0.0027995467343474847 0.0018842406149841935 -0.0009803697307184398 0.0015597988646485978 0.0033636031352137995 -0.0013679613503201132; 5.957832350481518e-17 -5.672925244624975e-17 1.658735815403105e-18 0.1542174992158436 0.21796528600528056 0.13198994379348716 0.13732278680715052 0.20278326844078884 0.14327901578660263 -0.011533536777490017 0.00012510516216442835 -0.018035643463844977 0.0029933913249113597 -0.0007189750167157034 -0.00028320470877725993 -0.13982816492846756 -0.011016167512193129 0.20498695205135647 -0.11286033157453505 0.023290779547916845 -0.011157162801290046 -0.0031740219537551646 0.0015951335013501735 -0.0010786050785834193 0.0017818286507094328 0.0035958648788115095 -0.0013124034798855972; 1.4884641559545154e-16 -1.9676320315043004e-17 8.285949303260346e-18 -0.0363619341993839 0.5027547143173895 0.20342078717952894 0.20278326844078884 0.45860698276131007 -0.020103252548307242 5.8795408909634115e-5 -0.0010059364015970649 -0.038845452322538075 0.0020954131429389723 -0.0013162473149686215 -1.630464751090788e-5 0.012966149376170218 0.003422950899523726 0.45713493283341006 -0.24623862877574676 0.05373623801529717 -0.0005685057569171245 -0.003349953614993612 0.009245121666156422 6.333967651118838e-5 -0.0012098669302554185 0.005635062862204953 -0.003207971635485146; -2.6158892480824602e-17 -1.5233218175986235e-16 -1.1446642093024695e-17 0.5447581101543878 -0.034929433164675616 0.12558158636649938 0.14327901578660263 -0.020103252548307242 0.4871917309988978 -0.036415280490110155 0.0018342650847155656 -0.0011987533456079198 0.006927487718327031 -0.00040694576959317985 -0.0008666689012223074 -0.46350306675139596 -0.039998410826840185 -0.010894421570093472 -0.003835952953275311 -0.004480362577419101 -0.03475337624397059 -0.00494703717405038 -0.008394806493437835 -0.003591901288186601 0.007541479254663582 0.0030882274883405826 0.0004124963438539906; 1.844872617411259e-18 1.1823574006427848e-17 -4.7139329501991025e-18 -0.04195510048875402 0.002186679578873358 -0.009884876441425145 -0.011533536777490017 5.8795408909634115e-5 -0.036415280490110155 0.003355149076834071 -0.00017113253378295021 2.7903992151544247e-5 -3.441674734288549e-5 2.7022719271761933e-6 8.23334362048528e-5 0.03734634470661608 0.003215379344001464 -0.0006006880883088634 -0.0012000376425958807 -0.0005729761675012921 0.0026932047881242827 0.0007557230034161155 0.0005723022804234376 0.0002064479395654276 -0.0004609463269569915 -0.0004786458544450129 -0.00010037076404187212; -4.196752724496389e-19 -5.416746242122389e-19 2.670186302440841e-19 0.002154513427850427 -0.0012182485546882798 3.9522459447686384e-5 0.00012510516216442835 -0.0010059364015970649 0.0018342650847155656 -0.00017113253378295021 1.0969491389705547e-5 8.536449843769953e-5 -7.617491925873939e-6 3.0533233723601022e-6 -4.179561820267352e-6 -0.0018902553523413874 -0.00016779210998946282 -0.0009705388530720606 0.000619358412169229 -8.22451404429099e-5 -0.00013202270512955975 -3.34352093248685e-5 -4.809585602317151e-5 -9.688370344533374e-6 2.434584754466226e-5 1.3418138796117441e-5 1.271735502231097e-5; -1.259491296360777e-17 2.476407050457178e-18 9.902981233405554e-19 0.00021478001431383942 -0.04284727947400572 -0.018079690582339526 -0.018035643463844977 -0.038845452322538075 -0.0011987533456079198 2.7903992151544247e-5 8.536449843769953e-5 0.0033617103084516995 -0.00036332091120805707 0.00012362283101896733 1.4567775860559774e-6 0.0008772554947334833 -0.00012000155953592911 -0.0387827695425048 0.021608353428790753 -0.004322252594433063 0.00022948948967266072 0.00020425120397100122 -0.0007274734458603151 3.4317880456748826e-5 2.850183693496739e-5 -0.0004298658615432112 0.00029344642080156655; 6.743962466892438e-19 -2.0234689214836416e-18 -4.451396093616671e-18 0.0066738119442117695 0.0030144217826222086 0.002995459962873844 0.0029933913249113597 0.0020954131429389723 0.006927487718327031 -3.441674734288549e-5 -7.617491925873939e-6 -0.00036332091120805707 0.0004941639530066145 -3.7642419236838695e-5 9.83276245306649e-7 -0.004528467546295607 -0.0003789513579290348 0.002240497708407808 -0.003103833739604801 -0.00036210405208010046 -0.0004349083338325164 0.0002037602401383398 -9.423977760866079e-5 -0.00010015560090120825 0.00018056549979491164 -0.00010540933000814182 -7.288495805116791e-5; -4.265016092358265e-19 1.8479027549073231e-19 2.6831269742696113e-19 -0.00034957855947544386 -0.001489378625558807 -0.0007204511116886534 -0.0007189750167157034 -0.0013162473149686215 -0.00040694576959317985 2.7022719271761933e-6 3.0533233723601022e-6 0.00012362283101896733 -3.7642419236838695e-5 5.842201489075393e-6 -7.3185988957165656e-9 0.0002689217829567287 1.6507208356302416e-5 -0.0013220482350060167 0.0008356415964272901 -0.00011468223490080425 3.0330022681610963e-5 -4.543629342934429e-6 -1.752483067883921e-5 6.41017211141022e-6 -8.754612541494913e-6 -7.701338907125194e-6 1.2988748628262737e-5; 3.9565163223496944e-20 2.8377258402809243e-19 -1.3389095381485578e-19 -0.0010029203595738237 3.7119631416326994e-5 -0.0002427072224056781 -0.00028320470877725993 -1.630464751090788e-5 -0.0008666689012223074 8.23334362048528e-5 -4.179561820267352e-6 1.4567775860559774e-6 9.83276245306649e-7 -7.3185988957165656e-9 2.028627727668926e-6 0.0008994175679549882 7.734873467333036e-5 -3.185987137512348e-5 -2.776826359703721e-5 -1.855947192077209e-5 6.451901991403527e-5 1.9569590635720513e-5 1.3171816753879684e-5 4.676237946656326e-6 -1.055000510990753e-5 -1.250067601447575e-5 -2.5725853773742526e-6; 2.4416066810159544e-17 1.467916473731728e-16 -1.2870466564366527e-17 -0.5235476858711073 0.03142164746853996 -0.12160105430170758 -0.13982816492846756 0.012966149376170218 -0.46350306675139596 0.03734634470661608 -0.0018902553523413874 0.0008772554947334833 -0.004528467546295607 0.0002689217829567287 0.0008994175679549882 0.4524920592502615 0.03901639405722845 0.004328085457333818 -0.002693331083320375 0.0003894137609299676 0.03347121771540934 0.00635305190654112 0.007751262480355128 0.0031526229510919533 -0.006736468317018984 -0.003995293570633398 -0.0006891036264413723; 2.8506203702269416e-18 1.2536739852398554e-17 -1.0696669620797387e-18 -0.04524950348585281 0.005234354521605073 -0.009443807843236254 -0.011016167512193129 0.003422950899523726 -0.039998410826840185 0.003215379344001464 -0.00016779210998946282 -0.00012000155953592911 -0.0003789513579290348 1.6507208356302416e-5 7.734873467333036e-5 0.03901639405722845 0.0033758056748421066 0.002672017376486892 -0.0014718722030148705 0.00030335572536341197 0.002878322663248411 0.0005302411253944728 0.0007136981835804876 0.00027165004853863363 -0.0005858828380685953 -0.00031570502696723525 -7.550304093452137e-5; 1.48013817644968e-16 -2.2479389061965568e-17 7.827037144299548e-18 -0.026083592007825382 0.5009452854191183 0.20530332455179046 0.20498695205135647 0.45713493283341006 -0.010894421570093472 -0.0006006880883088634 -0.0009705388530720606 -0.0387827695425048 0.002240497708407808 -0.0013220482350060167 -3.185987137512348e-5 0.004328085457333818 0.002672017376486892 0.4558405773823697 -0.24581475686641147 0.05349395994075671 -0.0012170285007858611 -0.0034198505647416856 0.009064121180014867 -6.867124799919116e-6 -0.001061140597848637 0.005670416074440992 -0.0031957769146251074; -8.060974977015126e-17 1.389303693462475e-17 1.5850258144832543e-17 0.007795253550627904 -0.2737210674981011 -0.11384336848933328 -0.11286033157453505 -0.24623862877574676 -0.003835952953275311 -0.0012000376425958807 0.000619358412169229 0.021608353428790753 -0.003103833739604801 0.0008356415964272901 -2.776826359703721e-5 -0.002693331083320375 -0.0014718722030148705 -0.24581475686641147 0.1408041163269181 -0.026127706921791578 0.0010301751304162348 0.0006025519931285303 -0.004626012214961456 0.0002982831322542953 6.458216060283465e-5 -0.0022990733753304768 0.0020008999972768336; 1.7108492080947968e-17 -2.0541067993702896e-18 7.536384440872769e-18 -0.005128486145404674 0.057605601051604054 0.023060030113865867 0.023290779547916845 0.05373623801529717 -0.004480362577419101 -0.0005729761675012921 -8.22451404429099e-5 -0.004322252594433063 -0.00036210405208010046 -0.00011468223490080425 -1.855947192077209e-5 0.0003894137609299676 0.00030335572536341197 0.05349395994075671 -0.026127706921791578 0.00717423199201098 -1.9542929705972215e-5 -0.000810717951765636 0.001150039959410032 9.632365482590683e-5 -0.0002919088535860381 0.0009156074370259354 -0.0002835148469118398; 1.2700948868758066e-18 1.1030517252671916e-17 -4.655899830429835e-21 -0.038980877842172516 0.00046287776926983907 -0.009848117628120642 -0.011157162801290046 -0.0005685057569171245 -0.03475337624397059 0.0026932047881242827 -0.00013202270512955975 0.00022948948967266072 -0.0004349083338325164 3.0330022681610963e-5 6.451901991403527e-5 0.03347121771540934 0.002878322663248411 -0.0012170285007858611 0.0010301751304162348 -1.9542929705972215e-5 0.002500526472935529 0.00042162927166863596 0.0005560613639668949 0.00024775409928402494 -0.000519852635712719 -0.00027809799468552106 -2.6764622706019242e-5; -6.696119442090515e-19 1.960313318173584e-18 -3.2494013161563927e-18 -0.006187195034468947 -0.002847784767902927 -0.0027995467343474847 -0.0031740219537551646 -0.003349953614993612 -0.00494703717405038 0.0007557230034161155 -3.34352093248685e-5 0.00020425120397100122 0.0002037602401383398 -4.543629342934429e-6 1.9569590635720513e-5 0.00635305190654112 0.0005302411253944728 -0.0034198505647416856 0.0006025519931285303 -0.000810717951765636 0.00042162927166863596 0.0003011900785325827 -2.7940256906295496e-20 -1.9463516485242953e-20 -1.0050505445851596e-5 -0.00021297235412297735 -2.6861257742656366e-5; 3.3078500592035365e-18 2.1048330148473913e-18 5.881564313829157e-19 -0.009601397385631415 0.010297009006599165 0.0018842406149841935 0.0015951335013501735 0.009245121666156422 -0.008394806493437835 0.0005723022804234376 -4.809585602317151e-5 -0.0007274734458603151 -9.423977760866079e-5 -1.752483067883921e-5 1.3171816753879684e-5 0.007751262480355128 0.0007136981835804876 0.009064121180014867 -0.004626012214961456 0.001150039959410032 0.0005560613639668949 -2.7940256906295496e-20 0.00031880138905730776 6.29277547091401e-5 -0.00015172363430819498 6.870310085160838e-5 -6.570497211201101e-5; 1.3022159717762597e-19 1.0926750149511406e-18 6.323587785736059e-19 -0.003886837103210349 4.993367362317162e-5 -0.0009803697307184398 -0.0010786050785834193 6.333967651118838e-5 -0.003591901288186601 0.0002064479395654276 -9.688370344533374e-6 3.4317880456748826e-5 -0.00010015560090120825 6.41017211141022e-6 4.676237946656326e-6 0.0031526229510919533 0.00027165004853863363 -6.867124799919116e-6 0.0002982831322542953 9.632365482590683e-5 0.00024775409928402494 -1.9463516485242953e-20 6.29277547091401e-5 3.267037759074216e-5 -6.533370952033481e-5 -1.0828657112765015e-6 5.437897753781058e-6; -6.384933325393985e-19 -2.2505531625316204e-18 -1.0736495926678668e-18 0.008256304443286854 -0.0013409779850553514 0.0015597988646485978 0.0017818286507094328 -0.0012098669302554185 0.007541479254663582 -0.0004609463269569915 2.434584754466226e-5 2.850183693496739e-5 0.00018056549979491164 -8.754612541494913e-6 -1.055000510990753e-5 -0.006736468317018984 -0.0005858828380685953 -0.001061140597848637 6.458216060283465e-5 -0.0002919088535860381 -0.000519852635712719 -1.0050505445851596e-5 -0.00015172363430819498 -6.533370952033481e-5 0.00013437880725617686 2.673526425991552e-22 1.1767314416274314e-20; 1.5643617507679333e-18 -1.4222930571743299e-18 2.051052415391445e-18 0.0037461756535838095 0.005666516559787287 0.0033636031352137995 0.0035958648788115095 0.005635062862204953 0.0030882274883405826 -0.0004786458544450129 1.3418138796117441e-5 -0.0004298658615432112 -0.00010540933000814182 -7.701338907125194e-6 -1.250067601447575e-5 -0.003995293570633398 -0.00031570502696723525 0.005670416074440992 -0.0022990733753304768 0.0009156074370259354 -0.00027809799468552106 -0.00021297235412297735 6.870310085160838e-5 -1.0828657112765015e-6 2.673526425991552e-22 0.00017598968195294928 -8.478098843039052e-6; -1.0971943724341637e-18 6.97275397205803e-21 6.378791949185355e-19 0.0007201717420216381 -0.003670482168746765 -0.0013679613503201132 -0.0013124034798855972 -0.003207971635485146 0.0004124963438539906 -0.00010037076404187212 1.271735502231097e-5 0.00029344642080156655 -7.288495805116791e-5 1.2988748628262737e-5 -2.5725853773742526e-6 -0.0006891036264413723 -7.550304093452137e-5 -0.0031957769146251074 0.0020008999972768336 -0.0002835148469118398 -2.6764622706019242e-5 -2.6861257742656366e-5 -6.570497211201101e-5 5.437897753781058e-6 1.1767314416274314e-20 -8.478098843039052e-6 3.2670377590742206e-5]

Note the use of mj_zeros to initialise the Jacobians, as outlined in Row vs. Column-Major Arrays.

Joint deviation cost

We will also want to add a cost on the joints deviating from the desired position qpos0. However, different parts of the humanoid will need to be treated differently. In particular:

  • The "free joints" (i.e: CoM position) should have no cost on them here, since we've already defined Qbalance to take care of them.
  • The joints in the left leg and lower abdomen should be held close to their original values to keep the balancing leg in place
  • Other joints, such as flailing limbs, should be allowed to move around more freely to maintain balance if the humanoid is perturbed.

Let's start by finding the array indices relevant to each set of joints. We'll leverage some of the nice features of the Named Access tools.

# Get indices into relevant sets of joints.
free_dofs = 1:6
body_dofs = 7:nv

# Get all the joints using a list comprehension.
# We add one to the raw ID to get the Julia 1-based index of the joint.
abdomen_dofs = [jnt.id+1 for jnt in MJ.joints(model) if occursin("abdomen", jnt.name)]
left_leg_dofs = [jnt.id+1 for jnt in MJ.joints(model) if occursin("left", jnt.name) && any(occursin(part, jnt.name) for part in ("hip", "knee", "ankle"))]

balance_dofs = vcat(abdomen_dofs, left_leg_dofs)
other_dofs = setdiff(body_dofs, balance_dofs)

println("Balance dofs: ", balance_dofs)
println("Other dofs:   ", other_dofs)
Balance dofs: [2, 3, 4, 11, 12, 13, 14, 15, 16]
Other dofs:   [7, 8, 9, 10, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27]

We can now use these indices to construct the second part of the $Q$ matrix.

# Cost coefficients
balance_joint_cost = 3          # Joints can move a bit and still balance
other_joint_cost   = 0.3        # Other joints can do whatever

# Construct joint Q matrix
Qjoint = Matrix{Float64}(I, nv, nv)
Qjoint[free_dofs, free_dofs] *= 0
Qjoint[balance_dofs, balance_dofs] *= balance_joint_cost
Qjoint[other_dofs, other_dofs] *= other_joint_cost

Putting this all together with Qbalance, we can construct our final $Q$ matrix defining our LQR cost.

balance_cost = 1000

Qpos = balance_cost*Qbalance + Qjoint
Q = [Qpos zeros(nv,nv); zeros(nv, 2nv)]  + (1e-10) * I # Add ϵI for positive definite Q
[1.0e-10 0.0 0.0 -3.4455790699220185e-14 1.6435007843370363e-13 6.068279684531847e-14 5.957832350481518e-14 1.4884641559545154e-13 -2.6158892480824602e-14 1.8448726174112593e-15 -4.1967527244963893e-16 -1.259491296360777e-14 6.743962466892437e-16 -4.265016092358265e-16 3.956516322349694e-17 2.4416066810159544e-14 2.8506203702269416e-15 1.4801381764496799e-13 -8.060974977015127e-14 1.7108492080947967e-14 1.2700948868758067e-15 -6.696119442090516e-16 3.3078500592035363e-15 1.3022159717762598e-16 -6.384933325393985e-16 1.5643617507679333e-15 -1.0971943724341636e-15 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0; 0.0 1.0e-10 0.0 -1.7026868264106963e-13 -1.6855275442690523e-14 -5.1006340013404474e-14 -5.672925244624975e-14 -1.9676320315043004e-14 -1.5233218175986235e-13 1.1823574006427848e-14 -5.416746242122389e-16 2.476407050457178e-15 -2.0234689214836414e-15 1.8479027549073231e-16 2.8377258402809245e-16 1.467916473731728e-13 1.2536739852398555e-14 -2.247938906196557e-14 1.389303693462475e-14 -2.0541067993702898e-15 1.1030517252671916e-14 1.9603133181735838e-15 2.1048330148473915e-15 1.0926750149511407e-15 -2.2505531625316204e-15 -1.42229305717433e-15 6.97275397205803e-18 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0; 0.0 0.0 1.0e-10 -1.740532938265174e-15 -8.265421392186514e-16 -7.983078331039035e-16 1.658735815403105e-15 8.285949303260345e-15 -1.1446642093024694e-14 -4.713932950199103e-15 2.6701863024408406e-16 9.902981233405554e-16 -4.451396093616671e-15 2.683126974269611e-16 -1.3389095381485577e-16 -1.2870466564366527e-14 -1.0696669620797387e-15 7.827037144299548e-15 1.5850258144832544e-14 7.536384440872768e-15 -4.655899830429835e-18 -3.249401316156393e-15 5.881564313829157e-16 6.323587785736059e-16 -1.0736495926678667e-15 2.051052415391445e-15 6.378791949185355e-16 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0; -3.4455790699220185e-14 -1.7026868264106963e-13 -1.740532938265174e-15 612.156694417341 -56.61728634672465 133.76853659664891 154.2174992158436 -36.361934199383896 544.7581101543879 -41.95510048875402 2.1545134278504268 0.21478001431383942 6.67381194421177 -0.34957855947544386 -1.0029203595738236 -523.5476858711073 -45.24950348585281 -26.08359200782538 7.795253550627904 -5.128486145404675 -38.98087784217252 -6.187195034468947 -9.601397385631415 -3.886837103210349 8.256304443286854 3.7461756535838093 0.720171742021638 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0; 1.6435007843370363e-13 -1.6855275442690523e-14 -8.265421392186514e-16 -56.61728634672465 553.6232120551902 219.73167831675926 217.96528600528057 502.75471431738947 -34.929433164675615 2.186679578873358 -1.2182485546882798 -42.84727947400572 3.0144217826222084 -1.489378625558807 0.037119631416326994 31.421647468539962 5.234354521605073 500.9452854191183 -273.7210674981011 57.605601051604054 0.4628777692698391 -2.847784767902927 10.297009006599165 0.049933673623171616 -1.3409779850553514 5.666516559787286 -3.670482168746765 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0; 6.068279684531847e-14 -5.1006340013404474e-14 -7.983078331039035e-16 133.76853659664891 219.73167831675926 127.46857263864614 131.98994379348716 203.42078717952893 125.58158636649938 -9.884876441425146 0.039522459447686384 -18.079690582339527 2.9954599628738436 -0.7204511116886534 -0.2427072224056781 -121.60105430170758 -9.443807843236254 205.30332455179047 -113.84336848933329 23.060030113865867 -9.848117628120642 -2.799546734347485 1.8842406149841935 -0.9803697307184398 1.5597988646485978 3.3636031352137996 -1.3679613503201131 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0; 5.957832350481518e-14 -5.672925244624975e-14 1.658735815403105e-15 154.2174992158436 217.96528600528057 131.98994379348716 137.62278680725052 202.78326844078885 143.27901578660263 -11.533536777490017 0.12510516216442835 -18.035643463844977 2.99339132491136 -0.7189750167157034 -0.28320470877725995 -139.82816492846754 -11.016167512193128 204.98695205135647 -112.86033157453505 23.290779547916845 -11.157162801290047 -3.174021953755165 1.5951335013501735 -1.0786050785834194 1.7818286507094327 3.5958648788115095 -1.3124034798855972 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0; 1.4884641559545154e-13 -1.9676320315043004e-14 8.285949303260345e-15 -36.361934199383896 502.75471431738947 203.42078717952893 202.78326844078885 458.90698276141006 -20.103252548307243 0.05879540890963411 -1.0059364015970649 -38.84545232253807 2.0954131429389724 -1.3162473149686216 -0.01630464751090788 12.966149376170218 3.422950899523726 457.13493283341006 -246.23862877574675 53.736238015297175 -0.5685057569171246 -3.349953614993612 9.245121666156422 0.06333967651118838 -1.2098669302554184 5.635062862204953 -3.207971635485146 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0; -2.6158892480824602e-14 -1.5233218175986235e-13 -1.1446642093024694e-14 544.7581101543879 -34.929433164675615 125.58158636649938 143.27901578660263 -20.103252548307243 487.4917309989978 -36.415280490110156 1.8342650847155655 -1.1987533456079198 6.927487718327031 -0.40694576959317985 -0.8666689012223074 -463.503066751396 -39.998410826840185 -10.894421570093472 -3.835952953275311 -4.480362577419101 -34.75337624397059 -4.947037174050379 -8.394806493437835 -3.5919012881866013 7.5414792546635825 3.0882274883405825 0.41249634385399064 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0; 1.8448726174112593e-15 1.1823574006427848e-14 -4.713932950199103e-15 -41.95510048875402 2.186679578873358 -9.884876441425146 -11.533536777490017 0.05879540890963411 -36.415280490110156 3.655149076934071 -0.1711325337829502 0.02790399215154425 -0.03441674734288549 0.0027022719271761935 0.0823334362048528 37.34634470661608 3.215379344001464 -0.6006880883088634 -1.2000376425958807 -0.5729761675012921 2.693204788124283 0.7557230034161155 0.5723022804234377 0.2064479395654276 -0.4609463269569915 -0.4786458544450129 -0.10037076404187212 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0; -4.1967527244963893e-16 -5.416746242122389e-16 2.6701863024408406e-16 2.1545134278504268 -1.2182485546882798 0.039522459447686384 0.12510516216442835 -1.0059364015970649 1.8342650847155655 -0.1711325337829502 3.0109694914897056 0.08536449843769953 -0.007617491925873939 0.003053323372360102 -0.004179561820267352 -1.8902553523413874 -0.1677921099894628 -0.9705388530720607 0.6193584121692289 -0.0822451404429099 -0.13202270512955974 -0.0334352093248685 -0.048095856023171506 -0.009688370344533374 0.02434584754466226 0.01341813879611744 0.01271735502231097 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0; -1.259491296360777e-14 2.476407050457178e-15 9.902981233405554e-16 0.21478001431383942 -42.84727947400572 -18.079690582339527 -18.035643463844977 -38.84545232253807 -1.1987533456079198 0.02790399215154425 0.08536449843769953 6.361710308551699 -0.3633209112080571 0.12362283101896733 0.0014567775860559774 0.8772554947334833 -0.12000155953592911 -38.7827695425048 21.608353428790753 -4.322252594433063 0.22948948967266072 0.20425120397100122 -0.7274734458603151 0.03431788045674883 0.028501836934967387 -0.4298658615432112 0.2934464208015665 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0; 6.743962466892437e-16 -2.0234689214836414e-15 -4.451396093616671e-15 6.67381194421177 3.0144217826222084 2.9954599628738436 2.99339132491136 2.0954131429389724 6.927487718327031 -0.03441674734288549 -0.007617491925873939 -0.3633209112080571 3.4941639531066144 -0.0376424192368387 0.000983276245306649 -4.528467546295607 -0.37895135792903484 2.240497708407808 -3.103833739604801 -0.36210405208010044 -0.43490833383251637 0.20376024013833982 -0.09423977760866079 -0.10015560090120824 0.18056549979491163 -0.10540933000814182 -0.07288495805116792 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0; -4.265016092358265e-16 1.8479027549073231e-16 2.683126974269611e-16 -0.34957855947544386 -1.489378625558807 -0.7204511116886534 -0.7189750167157034 -1.3162473149686216 -0.40694576959317985 0.0027022719271761935 0.003053323372360102 0.12362283101896733 -0.0376424192368387 3.0058422015890756 -7.318598895716566e-6 0.2689217829567287 0.016507208356302417 -1.3220482350060168 0.8356415964272901 -0.11468223490080424 0.030330022681610964 -0.004543629342934429 -0.017524830678839212 0.00641017211141022 -0.008754612541494913 -0.007701338907125194 0.012988748628262737 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0; 3.956516322349694e-17 2.8377258402809245e-16 -1.3389095381485577e-16 -1.0029203595738236 0.037119631416326994 -0.2427072224056781 -0.28320470877725995 -0.01630464751090788 -0.8666689012223074 0.0823334362048528 -0.004179561820267352 0.0014567775860559774 0.000983276245306649 -7.318598895716566e-6 3.002028627827669 0.8994175679549882 0.07734873467333035 -0.03185987137512348 -0.02776826359703721 -0.018559471920772088 0.06451901991403527 0.019569590635720512 0.013171816753879684 0.004676237946656326 -0.01055000510990753 -0.012500676014475751 -0.0025725853773742526 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0; 2.4416066810159544e-14 1.467916473731728e-13 -1.2870466564366527e-14 -523.5476858711073 31.421647468539962 -121.60105430170758 -139.82816492846754 12.966149376170218 -463.503066751396 37.34634470661608 -1.8902553523413874 0.8772554947334833 -4.528467546295607 0.2689217829567287 0.8994175679549882 455.4920592503615 39.016394057228446 4.3280854573338186 -2.693331083320375 0.3894137609299676 33.47121771540934 6.35305190654112 7.751262480355128 3.1526229510919532 -6.736468317018984 -3.9952935706333976 -0.6891036264413724 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0; 2.8506203702269416e-15 1.2536739852398555e-14 -1.0696669620797387e-15 -45.24950348585281 5.234354521605073 -9.443807843236254 -11.016167512193128 3.422950899523726 -39.998410826840185 3.215379344001464 -0.1677921099894628 -0.12000155953592911 -0.37895135792903484 0.016507208356302417 0.07734873467333035 39.016394057228446 3.6758056749421066 2.672017376486892 -1.4718722030148705 0.303355725363412 2.878322663248411 0.5302411253944728 0.7136981835804876 0.2716500485386336 -0.5858828380685953 -0.31570502696723524 -0.07550304093452137 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0; 1.4801381764496799e-13 -2.247938906196557e-14 7.827037144299548e-15 -26.08359200782538 500.9452854191183 205.30332455179047 204.98695205135647 457.13493283341006 -10.894421570093472 -0.6006880883088634 -0.9705388530720607 -38.7827695425048 2.240497708407808 -1.3220482350060168 -0.03185987137512348 4.3280854573338186 2.672017376486892 456.1405773824697 -245.81475686641147 53.49395994075671 -1.2170285007858612 -3.4198505647416857 9.064121180014867 -0.006867124799919116 -1.061140597848637 5.6704160744409915 -3.1957769146251076 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0; -8.060974977015127e-14 1.389303693462475e-14 1.5850258144832544e-14 7.795253550627904 -273.7210674981011 -113.84336848933329 -112.86033157453505 -246.23862877574675 -3.835952953275311 -1.2000376425958807 0.6193584121692289 21.608353428790753 -3.103833739604801 0.8356415964272901 -0.02776826359703721 -2.693331083320375 -1.4718722030148705 -245.81475686641147 141.10411632701812 -26.127706921791578 1.0301751304162348 0.6025519931285302 -4.626012214961456 0.2982831322542953 0.06458216060283464 -2.299073375330477 2.0008999972768335 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0; 1.7108492080947967e-14 -2.0541067993702898e-15 7.536384440872768e-15 -5.128486145404675 57.605601051604054 23.060030113865867 23.290779547916845 53.736238015297175 -4.480362577419101 -0.5729761675012921 -0.0822451404429099 -4.322252594433063 -0.36210405208010044 -0.11468223490080424 -0.018559471920772088 0.3894137609299676 0.303355725363412 53.49395994075671 -26.127706921791578 7.47423199211098 -0.019542929705972215 -0.810717951765636 1.150039959410032 0.09632365482590682 -0.29190885358603813 0.9156074370259354 -0.2835148469118398 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0; 1.2700948868758067e-15 1.1030517252671916e-14 -4.655899830429835e-18 -38.98087784217252 0.4628777692698391 -9.848117628120642 -11.157162801290047 -0.5685057569171246 -34.75337624397059 2.693204788124283 -0.13202270512955974 0.22948948967266072 -0.43490833383251637 0.030330022681610964 0.06451901991403527 33.47121771540934 2.878322663248411 -1.2170285007858612 1.0301751304162348 -0.019542929705972215 2.800526473035529 0.42162927166863595 0.5560613639668949 0.24775409928402495 -0.519852635712719 -0.27809799468552104 -0.02676462270601924 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0; -6.696119442090516e-16 1.9603133181735838e-15 -3.249401316156393e-15 -6.187195034468947 -2.847784767902927 -2.799546734347485 -3.174021953755165 -3.349953614993612 -4.947037174050379 0.7557230034161155 -0.0334352093248685 0.20425120397100122 0.20376024013833982 -0.004543629342934429 0.019569590635720512 6.35305190654112 0.5302411253944728 -3.4198505647416857 0.6025519931285302 -0.810717951765636 0.42162927166863595 0.6011900786325827 -2.7940256906295493e-17 -1.9463516485242952e-17 -0.010050505445851596 -0.21297235412297735 -0.026861257742656366 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0; 3.3078500592035363e-15 2.1048330148473915e-15 5.881564313829157e-16 -9.601397385631415 10.297009006599165 1.8842406149841935 1.5951335013501735 9.245121666156422 -8.394806493437835 0.5723022804234377 -0.048095856023171506 -0.7274734458603151 -0.09423977760866079 -0.017524830678839212 0.013171816753879684 7.751262480355128 0.7136981835804876 9.064121180014867 -4.626012214961456 1.150039959410032 0.5560613639668949 -2.7940256906295493e-17 0.6188013891573078 0.0629277547091401 -0.151723634308195 0.06870310085160838 -0.06570497211201101 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0; 1.3022159717762598e-16 1.0926750149511407e-15 6.323587785736059e-16 -3.886837103210349 0.049933673623171616 -0.9803697307184398 -1.0786050785834194 0.06333967651118838 -3.5919012881866013 0.2064479395654276 -0.009688370344533374 0.03431788045674883 -0.10015560090120824 0.00641017211141022 0.004676237946656326 3.1526229510919532 0.2716500485386336 -0.006867124799919116 0.2982831322542953 0.09632365482590682 0.24775409928402495 -1.9463516485242952e-17 0.0629277547091401 0.33267037769074215 -0.0653337095203348 -0.0010828657112765016 0.005437897753781058 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0; -6.384933325393985e-16 -2.2505531625316204e-15 -1.0736495926678667e-15 8.256304443286854 -1.3409779850553514 1.5597988646485978 1.7818286507094327 -1.2098669302554184 7.5414792546635825 -0.4609463269569915 0.02434584754466226 0.028501836934967387 0.18056549979491163 -0.008754612541494913 -0.01055000510990753 -6.736468317018984 -0.5858828380685953 -1.061140597848637 0.06458216060283464 -0.29190885358603813 -0.519852635712719 -0.010050505445851596 -0.151723634308195 -0.0653337095203348 0.43437880735617684 2.673526425991552e-19 1.1767314416274314e-17 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0; 1.5643617507679333e-15 -1.42229305717433e-15 2.051052415391445e-15 3.7461756535838093 5.666516559787286 3.3636031352137996 3.5958648788115095 5.635062862204953 3.0882274883405825 -0.4786458544450129 0.01341813879611744 -0.4298658615432112 -0.10540933000814182 -0.007701338907125194 -0.012500676014475751 -3.9952935706333976 -0.31570502696723524 5.6704160744409915 -2.299073375330477 0.9156074370259354 -0.27809799468552104 -0.21297235412297735 0.06870310085160838 -0.0010828657112765016 2.673526425991552e-19 0.4759896820529493 -0.008478098843039052 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0; -1.0971943724341636e-15 6.97275397205803e-18 6.378791949185355e-16 0.720171742021638 -3.670482168746765 -1.3679613503201131 -1.3124034798855972 -3.207971635485146 0.41249634385399064 -0.10037076404187212 0.01271735502231097 0.2934464208015665 -0.07288495805116792 0.012988748628262737 -0.0025725853773742526 -0.6891036264413724 -0.07550304093452137 -3.1957769146251076 2.0008999972768335 -0.2835148469118398 -0.02676462270601924 -0.026861257742656366 -0.06570497211201101 0.005437897753781058 1.1767314416274314e-17 -0.008478098843039052 0.3326703776907422 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0; 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 1.0e-10 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0; 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 1.0e-10 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0; 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 1.0e-10 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0; 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 1.0e-10 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0; 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 1.0e-10 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0; 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 1.0e-10 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0; 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 1.0e-10 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0; 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 1.0e-10 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0; 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 1.0e-10 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0; 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 1.0e-10 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0; 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 1.0e-10 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0; 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 1.0e-10 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0; 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 1.0e-10 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0; 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 1.0e-10 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0; 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 1.0e-10 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0; 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 1.0e-10 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0; 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 1.0e-10 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0; 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 1.0e-10 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0; 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 1.0e-10 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0; 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 1.0e-10 0.0 0.0 0.0 0.0 0.0 0.0 0.0; 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 1.0e-10 0.0 0.0 0.0 0.0 0.0 0.0; 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 1.0e-10 0.0 0.0 0.0 0.0 0.0; 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 1.0e-10 0.0 0.0 0.0 0.0; 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 1.0e-10 0.0 0.0 0.0; 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 1.0e-10 0.0 0.0; 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 1.0e-10 0.0; 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 1.0e-10]

Note that the balance_cost is quite large in comparison to balance_joint_cost because the units for the CoM position (metres) are typically "larger" than the units for joint angles (radians). Eg: if the CoM shifts by $0.1\,$m, the humanoid will most likely fall over, but a $0.1\,$radian change in a leg angle will probably be fine.

Computing optimal control gains

With both cost matrices defined, computing the optimal control gains is simply a matter of linearising the system dynamics about our chosen set-point (qpos0, ctrl0) and solving the LQR Riccati equations. We can follow the same process as our Balancing a Cart-Pole example. Let's start with the linear system Jacobians.

# Initialise the model at our set point
reset!(model, data)
data.ctrl .= ctrl0
data.qpos .= qpos0

# Finite-difference parameters
ϵ = 1e-6
centred = true

# Compute the Jacobians
A = mj_zeros(2nv, 2nv)
B = mj_zeros(2nv, nu)
mjd_transitionFD(model, data, ϵ, centred, A, B, nothing, nothing)
@show A, B
([0.999999999999965 -1.6728900708272432e-14 … 5.182409077901198e-7 -1.2636514314691156e-6; 5.929230630780102e-15 1.0000000000000053 … -1.0814096869704907e-6 3.184198256223113e-7; … ; 9.6710833785707e-11 1.5178830414797062e-11 … 0.9863941045039348 0.005423895092608599; -8.673617379884035e-12 1.3010426069826053e-11 … 0.005421714875664119 0.9816818517553381], [7.462459666320812e-6 0.00015992280042538722 … -5.1824089265999375e-5 0.0002527302842909079; -6.566536143877671e-5 -1.892038546777615e-5 … 0.00010814097135673252 -6.368396561065917e-5; … ; -0.9215383986256771 -0.49972383692198363 … 1.3605895558448446 -1.0847789970640578; 0.5151568750298707 -0.05289604147230842 … -0.5421714840926282 3.663629652986104])

Just like the cart-pole example, we'll directly use ared from MatrixEquations.jl to compute the LQR gain matrix $K$ as a lightweight alternative to loading ControlSystemsCore.jl.

using MatrixEquations

S = zeros(2nv, nu)
_, _, K, _ = ared(A,B,R,Q,S)
([0.020288540546339554 -0.008811423080995955 … 5.737936700028865e-5 -1.6412776321650778e-5; -0.008811423080995955 0.08914012180232787 … 2.699166768109102e-5 -2.565529434759259e-5; … ; 5.737936700028865e-5 2.699166768109102e-5 … 0.3155587930622829 -0.009967695851371685; -1.6412776321650778e-5 -2.565529434759259e-5 … -0.009967695851371685 0.11841745479161118], ComplexF64[0.4312302153190349 - 0.0im, 0.4471490710952473 - 0.0im, 0.6008161459180423 - 0.06411506769387793im, 0.6008161459180423 + 0.06411506769387793im, 0.6451316918216 - 0.0im, 0.6552885878168582 - 0.1575204218594505im, 0.6552885878168582 + 0.1575204218594505im, 0.672209240193633 - 0.16833172875992716im, 0.672209240193633 + 0.16833172875992716im, 0.7171920017054538 - 0.0im  …  0.9806133270196457 + 0.024405600039142522im, 0.9823812632183481 - 0.01017312446598192im, 0.9823812632183481 + 0.01017312446598192im, 0.9852961181696518 - 0.008929029861804284im, 0.9852961181696518 + 0.008929029861804284im, 0.9922596092180211 - 0.011321177180149883im, 0.9922596092180211 + 0.011321177180149883im, 0.9999122434225632 - 0.0im, 0.999999994783954 - 0.0im, 0.9999999988971865 - 0.0im], [3.7048390161276146e-6 4.983602616581706e-6 … -0.005981715973722129 0.018629039546679084; 7.762126581711826e-6 -2.515026011451482e-6 … -0.004829461294098559 0.0017470454566441331; … ; -1.0023933755820558e-6 -2.609762684846473e-7 … 0.06317369414167193 0.007868312338494846; -4.529837036582052e-7 -3.5095142036922766e-7 … 0.011202192710459637 0.036556498995066], [2.5241064396538004e-6 -3.101447355031313e-7 … 0.9915800562107996 0.1246370224097896; -4.646651813690919e-6 3.484677205177839e-6 … 0.12511894232908122 -0.9901041113206177; … ; -1.111169912826865e-6 -2.8647703936078827e-6 … 0.00023778847491074594 -0.00040297162050049323; 9.257813427651435e-7 4.187571346626105e-6 … 0.0008293709294834153 0.00016616919556029092], (Sx = LinearAlgebra.UniformScaling{Float64}(1.2353087992814715), Sxi = LinearAlgebra.UniformScaling{Float64}(1.2353087992814715), Sr = LinearAlgebra.UniformScaling{Float64}(0.8095141883403235)))

Testing the controller

We are now ready to test our controller. First, we'll write a function to set the control gains on our humanoid at each simulation time-step. Note the use of mj_differentiatePos to compute the position error, which uses a finite difference to calculate the difference between two positions rather than subtracting qpos0 - data.qpos. This is necessary because orientations of free bodies in MuJoCo are represented by 4-element quaternions, but we are interested in differences between points in 3D space. This is why model.nq >= model.nv for all MuJoCo models.

function humanoid_ctrl!(m::Model, d::Data)
    Δq = zeros(m.nv)
    mj_differentiatePos(m, Δq, 1, qpos0, d.qpos)

    Δx = vcat(Δq, data.qvel)
    data.ctrl .= ctrl0 .- K*Δx
    nothing
end
Performance Tip

This function captures non-const global variables and will take a performance hit (see Performance Tips) for more details. To remedy the performance hit, one can use functors instead, as described below.

Let's run the visualiser and test our controller by manually perturbing the humanoid. You can do this by double-clicking on a body to select it, then CTRL+RightClick and drag to apply a force.

reset!(model, data)
data.qpos .= qpos0
visualise!(model, data, controller=humanoid_ctrl!)

For small perturbations, our controller works just great! If you apply a large perturbation, you'll find that it's easy to de-stabilise the system and get the humanoid to fall over. This is to be expected from a simple LQR, but feel free to play around with the weight matrices to see if you can improve the robustness. In general, designing robust controllers for nonlinear systems is a difficult challenge and an open area of research in the control community.

Improve performance with functors and cached memory

A functor in Julia can be created by defining a struct to hold any variables that may be needed by a function, and defining the function itself when the struct is called. The struct will store all the variables that are used within the function that are not accessible from the input arguments. Our humanoid controller requires qpos0, ctrl0 and K. We also can reuse some memory in the form of Δq and Δx, so we can store this too inside the functor. Finally, we can store a matrix which holds the output of the K*Δx matrix multiplication.

struct LQRHumanoidBalance{TQ,TC,TK,TDQ,TDX,TKDX}
    qpos0::TQ
    ctrl0::TC
    K::TK
    Δq::TDQ
    Δx::TDX
    KΔx::TKDX
end

Notice that we are not restricting the types of the fields for clarity, however, the true type will be stored in the functors type information.

Next, we must define the actual function body for the functor, which follows the same format as the previous controller:

function (functor::LQRHumanoidBalance)(model::Model, data::Data)
    Δq = functor.Δq
    mj_differentiatePos(model, Δq, 1, functor.qpos0, data.qpos)

    Δx = functor.Δx
    Δx[1:length(Δq)] .= Δq
    Δx[(length(Δq)+1):end] .= data.qvel

    KΔx = functor.KΔx
    mul!(KΔx, functor.K, Δx)

    data.ctrl .= functor.ctrl0 .- KΔx
    nothing
end

We can now create an instance of the functor, creating memory for the caches as well:

controller = LQRHumanoidBalance(
    qpos0,
    ctrl0,
    K,
    zeros(eltype(qpos0), model.nv), # Δq
    zeros(eltype(qpos0), 2 * model.nv), # Δx
    similar(ctrl0) # KΔx
);

Now we can benchmark these two functions with BenchmarkTools.jl:

using BenchmarkTools
@benchmark humanoid_ctrl!(model, data)
BenchmarkTools.Trial: 10000 samples with 10 evaluations.
 Range (minmax):  1.160 μs  4.426 μs   GC (min … max): 0.00% … 0.00%
 Time  (median):     1.664 μs                GC (median):    0.00%
 Time  (mean ± σ):   1.613 μs ± 259.292 ns   GC (mean ± σ):  0.00% ± 0.00%

                     █▆       ▇▂                              
  ▄█▇▅▄▂▂▁▁▃▆▄▂▁▁▁▁▁▅███▄▂▂▂▆██▆▄▂▂▂▄▇▇▆▃▂▂▁▂▂▃▃▂▂▁▁▁▁▁▁▁▁▁ ▃
  1.16 μs         Histogram: frequency by time        2.23 μs <

 Memory estimate: 1.22 KiB, allocs estimate: 10.
@benchmark controller(model, data)
BenchmarkTools.Trial: 10000 samples with 102 evaluations.
 Range (minmax):  784.020 ns 1.338 μs   GC (min … max): 0.00% … 0.00%
 Time  (median):     787.657 ns               GC (median):    0.00%
 Time  (mean ± σ):   797.755 ns ± 49.732 ns   GC (mean ± σ):  0.00% ± 0.00%

  █            ▁                                             ▁
  ██▇▄▃▁▄▁▄▁▄▆█▇▆▅▆▆▅▅▆▅▄▄▄▁▄▄▄▄▁▄▁▃▃▃▁▃▄▄▅▅▄▃▅▄▆▆▆▆▄▆▅▆▇▆▅ █
  784 ns        Histogram: log(frequency) by time      1.09 μs <

 Memory estimate: 0 bytes, allocs estimate: 0.

We can see that using a functor will give the best performance, at the cost of slightly more boilerplate code.

Thread Safety

The use of the functor with local cache is not thread-safe, and a different functor must be used on each thread if running in parallel.