Skip to content

Commit 78c2460

Browse files
lvjonoklvjonok
authored
Feature/energy match (#4)
* compare energy * feat: matched skydio here * feat: complete test of energy and parameters
1 parent d3cd436 commit 78c2460

File tree

6 files changed

+471
-20
lines changed

6 files changed

+471
-20
lines changed

‎examples/match_energy.ipynb‎

Lines changed: 240 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,240 @@
1+
{
2+
"cells": [
3+
{
4+
"cell_type": "code",
5+
"execution_count": 1,
6+
"metadata": {},
7+
"outputs": [],
8+
"source": [
9+
"import mujoco\n",
10+
"import numpy as np\n",
11+
"import pinocchio as pin\n",
12+
"\n",
13+
"# from robot_descriptions.skydio_x2_description import URDF_PATH\n",
14+
"# from robot_descriptions.skydio_x2_mj_description import MJCF_PATH\n",
15+
"from robot_descriptions.z1_description import URDF_PATH\n",
16+
"from robot_descriptions.z1_mj_description import MJCF_PATH\n",
17+
"\n",
18+
"from mujoco_sysid import parameters, regressors\n",
19+
"from mujoco_sysid.utils import muj2pin"
20+
]
21+
},
22+
{
23+
"cell_type": "code",
24+
"execution_count": 2,
25+
"metadata": {},
26+
"outputs": [],
27+
"source": [
28+
"mjmodel = mujoco.MjModel.from_xml_path(MJCF_PATH)\n",
29+
"# enable energy\n",
30+
"mjmodel.opt.enableflags |= mujoco.mjtEnableBit.mjENBL_ENERGY\n",
31+
"# disable friction, contact and limits\n",
32+
"mjmodel.opt.disableflags |= (\n",
33+
" mujoco.mjtDisableBit.mjDSBL_CONTACT | mujoco.mjtDisableBit.mjDSBL_FRICTIONLOSS | mujoco.mjtDisableBit.mjDSBL_LIMIT\n",
34+
")\n",
35+
"mjdata = mujoco.MjData(mjmodel)\n",
36+
"\n",
37+
"pinmodel = pin.buildModelFromUrdf(URDF_PATH)\n",
38+
"pindata = pin.Data(pinmodel)"
39+
]
40+
},
41+
{
42+
"cell_type": "code",
43+
"execution_count": 3,
44+
"metadata": {},
45+
"outputs": [],
46+
"source": [
47+
"# np.random.seed(10)"
48+
]
49+
},
50+
{
51+
"cell_type": "code",
52+
"execution_count": 4,
53+
"metadata": {},
54+
"outputs": [],
55+
"source": [
56+
"q, v = np.random.randn(mjmodel.nq), np.zeros(mjmodel.nv)\n",
57+
"\n",
58+
"\n",
59+
"pinq, pinv = muj2pin(q, v)"
60+
]
61+
},
62+
{
63+
"cell_type": "code",
64+
"execution_count": 5,
65+
"metadata": {},
66+
"outputs": [
67+
{
68+
"data": {
69+
"text/plain": [
70+
"(array([-1.21974604, 0. ]), -1.2197460404740708)"
71+
]
72+
},
73+
"execution_count": 5,
74+
"metadata": {},
75+
"output_type": "execute_result"
76+
}
77+
],
78+
"source": [
79+
"mjdata.qpos[:] = q.copy()\n",
80+
"mjdata.qvel[:] = v.copy()\n",
81+
"\n",
82+
"mujoco.mj_step(mjmodel, mjdata)\n",
83+
"\n",
84+
"mj_en = mjdata.energy.copy()\n",
85+
"mj_en[0] += regressors.potential_energy_bias(mjmodel)\n",
86+
"\n",
87+
"mj_en, np.sum(mj_en)"
88+
]
89+
},
90+
{
91+
"cell_type": "code",
92+
"execution_count": 6,
93+
"metadata": {},
94+
"outputs": [
95+
{
96+
"data": {
97+
"text/plain": [
98+
"(-1.2197457679147197, [-1.2197457679147197, 0.0])"
99+
]
100+
},
101+
"execution_count": 6,
102+
"metadata": {},
103+
"output_type": "execute_result"
104+
}
105+
],
106+
"source": [
107+
"(\n",
108+
" np.sum(\n",
109+
" [\n",
110+
" pin.computePotentialEnergy(pinmodel, pindata, pinq),\n",
111+
" pin.computeKineticEnergy(pinmodel, pindata, pinq, pinv),\n",
112+
" ]\n",
113+
" ),\n",
114+
" [\n",
115+
" pin.computePotentialEnergy(pinmodel, pindata, pinq),\n",
116+
" pin.computeKineticEnergy(pinmodel, pindata, pinq, pinv),\n",
117+
" ],\n",
118+
")"
119+
]
120+
},
121+
{
122+
"cell_type": "code",
123+
"execution_count": 7,
124+
"metadata": {},
125+
"outputs": [
126+
{
127+
"data": {
128+
"text/plain": [
129+
"((60,),\n",
130+
" array([ 6.73326000e-01, 1.66311522e-06, -1.69664685e-04, 1.56021081e-02,\n",
131+
" 1.64484975e-03, -5.95801269e-08, 1.08083746e-03, -4.38536417e-07,\n",
132+
" 4.43147004e-06, 8.39403034e-04]))"
133+
]
134+
},
135+
"execution_count": 7,
136+
"metadata": {},
137+
"output_type": "execute_result"
138+
}
139+
],
140+
"source": [
141+
"theta = np.concatenate([parameters.get_dynamic_parameters(mjmodel, i) for i in mjmodel.jnt_bodyid])\n",
142+
"\n",
143+
"theta.shape, theta[:10]"
144+
]
145+
},
146+
{
147+
"cell_type": "code",
148+
"execution_count": 8,
149+
"metadata": {},
150+
"outputs": [
151+
{
152+
"name": "stdout",
153+
"output_type": "stream",
154+
"text": [
155+
"for body 0 norm of difference is 4.903343079507622e-07\n",
156+
"for body 1 norm of difference is 2.5970031568529883e-06\n",
157+
"for body 2 norm of difference is 2.650078190181762e-07\n",
158+
"for body 3 norm of difference is 3.701542518534774e-07\n",
159+
"for body 4 norm of difference is 8.053173890678282e-08\n",
160+
"for body 5 norm of difference is 7.002258930267969e-08\n"
161+
]
162+
},
163+
{
164+
"data": {
165+
"text/plain": [
166+
"((60,), 2.6839308799870235e-06)"
167+
]
168+
},
169+
"execution_count": 8,
170+
"metadata": {},
171+
"output_type": "execute_result"
172+
}
173+
],
174+
"source": [
175+
"params = []\n",
176+
"\n",
177+
"for i in range(len(pinmodel.inertias) - 1):\n",
178+
" params.extend(pinmodel.inertias[i + 1].toDynamicParameters())\n",
179+
"\n",
180+
" last_params = np.array(params[-10:])\n",
181+
" last_theta = theta[i * 10 : (i + 1) * 10]\n",
182+
"\n",
183+
" # mass should match\n",
184+
" assert np.isclose(last_params[0], last_theta[0])\n",
185+
"\n",
186+
" # lever arm should match\n",
187+
" assert np.allclose(last_params[1:4], last_theta[1:4])\n",
188+
"\n",
189+
" print(f\"for body {i} norm of difference is {np.linalg.norm(last_params - last_theta)}\")\n",
190+
"\n",
191+
"params = np.array(params)\n",
192+
"\n",
193+
"params.shape, np.linalg.norm(params - theta)"
194+
]
195+
},
196+
{
197+
"cell_type": "code",
198+
"execution_count": 9,
199+
"metadata": {},
200+
"outputs": [
201+
{
202+
"data": {
203+
"text/plain": [
204+
"-1.2197460404740712"
205+
]
206+
},
207+
"execution_count": 9,
208+
"metadata": {},
209+
"output_type": "execute_result"
210+
}
211+
],
212+
"source": [
213+
"reg_en = regressors.mj_energyRegressor(mjmodel, mjdata)[2]\n",
214+
"\n",
215+
"reg_en @ theta"
216+
]
217+
}
218+
],
219+
"metadata": {
220+
"kernelspec": {
221+
"display_name": "venv",
222+
"language": "python",
223+
"name": "python3"
224+
},
225+
"language_info": {
226+
"codemirror_mode": {
227+
"name": "ipython",
228+
"version": 3
229+
},
230+
"file_extension": ".py",
231+
"mimetype": "text/x-python",
232+
"name": "python",
233+
"nbconvert_exporter": "python",
234+
"pygments_lexer": "ipython3",
235+
"version": "3.10.12"
236+
}
237+
},
238+
"nbformat": 4,
239+
"nbformat_minor": 2
240+
}

‎mujoco_sysid/parameters.py‎

Lines changed: 25 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,7 @@ def get_dynamic_parameters(mjmodel, body_id) -> npt.ArrayLike:
1818
Returns:
1919
npt.ArrayLike: theta of the body
2020
"""
21-
mass = mjmodel.body(body_id).mass
21+
mass = mjmodel.body(body_id).mass[0]
2222
rc = mjmodel.body(body_id).ipos
2323
diag_inertia = mjmodel.body(body_id).inertia
2424

@@ -28,10 +28,21 @@ def get_dynamic_parameters(mjmodel, body_id) -> npt.ArrayLike:
2828

2929
R = r_flat.reshape(3, 3)
3030

31-
shift = -mass * skew(rc) @ skew(rc)
32-
mjinertia = R.T @ np.diag(diag_inertia) @ R + shift
31+
shift = mass * skew(rc) @ skew(rc)
32+
mjinertia = R @ np.diag(diag_inertia) @ R.T - shift
3333

34-
return np.concatenate([mass, mass * rc, mjinertia[np.triu_indices(3)]])
34+
upper_triangular = np.array(
35+
[
36+
mjinertia[0, 0],
37+
mjinertia[0, 1],
38+
mjinertia[1, 1],
39+
mjinertia[0, 2],
40+
mjinertia[1, 2],
41+
mjinertia[2, 2],
42+
]
43+
)
44+
45+
return np.concatenate([[mass], mass * rc, upper_triangular])
3546

3647

3748
def set_dynamic_parameters(mjmodel, body_id, theta: npt.ArrayLike) -> None:
@@ -46,14 +57,19 @@ def set_dynamic_parameters(mjmodel, body_id, theta: npt.ArrayLike) -> None:
4657
mass = theta[0]
4758
rc = theta[1:4] / mass
4859
inertia = theta[4:]
49-
inertia_full = np.zeros((3, 3))
50-
inertia_full[np.triu_indices(3)] = inertia
60+
inertia_full = np.array(
61+
[
62+
[inertia[0], inertia[1], inertia[3]],
63+
[inertia[1], inertia[2], inertia[4]],
64+
[inertia[3], inertia[4], inertia[5]],
65+
]
66+
)
5167

5268
# shift the inertia
53-
inertia_full -= -mass * skew(rc) @ skew(rc)
69+
inertia_full += mass * skew(rc) @ skew(rc)
5470

5571
# eigen decomposition
56-
eigval, eigvec = np.linalg.eig(inertia_full)
72+
eigval, eigvec = np.linalg.eigh(inertia_full)
5773
R = eigvec
5874
diag_inertia = eigval
5975

@@ -62,7 +78,7 @@ def set_dynamic_parameters(mjmodel, body_id, theta: npt.ArrayLike) -> None:
6278
raise ValueError("Cannot deduce inertia matrix because RIR^T is singular.")
6379

6480
# set the mass
65-
mjmodel.body(body_id).mass = mass
81+
mjmodel.body(body_id).mass = np.array([mass])
6682
mjmodel.body(body_id).ipos = rc
6783

6884
# set the orientation

‎mujoco_sysid/regressors.py‎

Lines changed: 31 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -233,7 +233,16 @@ def body_energyRegressor(
233233
return kinetic, potential
234234

235235

236-
def mj_energyRegressor(mj_model, mj_data, body_offset: int = 0) -> tuple[npt.ArrayLike, npt.ArrayLike, npt.ArrayLike]:
236+
def get_jacobian(mjmodel, mjdata, bodyid):
237+
R = mjdata.xmat[bodyid].reshape(3, 3)
238+
239+
jacp, jacr = np.zeros((3, 6)), np.zeros((3, 6))
240+
mujoco.mj_jacBody(mjmodel, mjdata, jacp, jacr, bodyid)
241+
242+
return np.vstack([R.T @ jacr, R.T @ jacp])
243+
244+
245+
def mj_energyRegressor(mj_model, mj_data) -> tuple[npt.ArrayLike, npt.ArrayLike, npt.ArrayLike]:
237246
"""
238247
mj_energyRegressor returns kinetic, potential, and total energy regressors for the whole model.
239248
@@ -263,15 +272,31 @@ def mj_energyRegressor(mj_model, mj_data, body_offset: int = 0) -> tuple[npt.Arr
263272
potential_regressor = np.zeros(njoints * 10)
264273
velocity = np.zeros(6)
265274

266-
for i in range(njoints):
267-
mujoco.mj_objectVelocity(mj_model, mj_data, 2, i + body_offset, velocity, 1)
268-
rotation = mj_data.xmat[i + body_offset].reshape(3, 3).copy()
269-
position = mj_data.xpos[i + body_offset]
270-
275+
for i, bodyid in enumerate(mj_model.jnt_bodyid):
276+
# same as jacobian @ qvel
277+
mujoco.mj_objectVelocity(mj_model, mj_data, 2, bodyid, velocity, 1)
271278
v, w = velocity[3:], velocity[:3]
279+
280+
rotation = mj_data.xmat[bodyid].reshape(3, 3)
281+
position = mj_data.xpos[bodyid]
282+
272283
kinetic, potential = body_energyRegressor(v, w, position, rotation)
273284
kinetic_regressor[10 * i : 10 * (i + 1)] = kinetic
274285
potential_regressor[10 * i : 10 * (i + 1)] = potential
275286
energy_regressor[10 * i : 10 * (i + 1)] = kinetic + potential
276287

277288
return kinetic_regressor, potential_regressor, energy_regressor
289+
290+
291+
def potential_energy_bias(mjmodel):
292+
"""
293+
The bodies before the first joint are considered to be fixed in space.
294+
They are included in potential energy calculation, but not in the regressor.
295+
"""
296+
297+
bias = 0
298+
for i in range(mjmodel.nbody):
299+
if i not in mjmodel.jnt_bodyid:
300+
bias += mjmodel.body(i).mass[0] * mjmodel.opt.gravity[2] * mjmodel.body(i).ipos[2]
301+
302+
return bias

0 commit comments

Comments
 (0)