chenlin/Mujoco_WASM
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Merge pull request #1450 from v-r-a:angmomMat

Browse Source 
PiperOrigin-RevId: 612835207
Change-Id: I8e7b01cc433563556a3ae4bbdd0c53612cc15e97
This commit is contained in:
Copybara-Service
2024-03-05 07:35:15 -08:00
parent 49ddb7ca2e 93f9154d44
commit 6b7d714291
12 changed files with 279 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files
doc/APIreference/functions.rst
+12
View File
@@ -307,6 +307,18 @@ mj_jacPointAxis
Compute translation end-effector Jacobian of point, and rotation Jacobian of axis.
.. _mj_angmomMat:
mj_angmomMat
~~~~~~~~~~~~
.. mujoco-include:: mj_angmomMat
This function computes the ``3 x nv`` angular momentum matrix :math:`H(q)`, providing the linear mapping from
generalized velocities to subtree angular momentum. More precisely if :math:`h` is the subtree angular momentum of
body index ``body`` in ``mjData.subtree_angmom`` (reported by the :ref:`subtreeangmom<sensor-subtreeangmom>` sensor)
and :math:`\dot q` is the generalized velocity ``mjData.qvel``, then :math:`h = H \dot q`.
.. _mj_name2id:
mj_name2id
doc/APIreference/functions_override.rst
+7
View File
@@ -173,6 +173,13 @@ the advantages of working in minimal coordinates.
This and the remaining variants of the Jacobian function call mj_jac internally, with the center of the body, geom or
site. They are just shortcuts; the same can be achieved by calling mj_jac directly.
.. _mj_angmomMat:
This function computes the ``3 x nv`` angular momentum matrix :math:`H(q)`, providing the linear mapping from
generalized velocities to subtree angular momentum. More precisely if :math:`h` is the subtree angular momentum of
body index ``body`` in ``mjData.subtree_angmom`` (reported by the :ref:`subtreeangmom<sensor-subtreeangmom>` sensor)
and :math:`\dot q` is the generalized velocity ``mjData.qvel``, then :math:`h = H \dot q`.
.. _mj_mulM:
This function multiplies the joint-space inertia matrix stored in mjData.qM by a vector. qM has a custom sparse format
doc/changelog.rst
+11 -8
View File
@@ -16,27 +16,30 @@ General
its internal file buffers. :ref:`mj_addBufferVFS` allocates an empty buffer with a given name in an mjVFS and
copies the data buffer into it, combining and replacing the deprecated two-step process of calling
:ref:`mj_makeEmptyFileVFS` followed by a direct copy into the given mjVFS internal file buffer.
3. Added :ref:`mj_angmomMat` which computes the ``3 x nv`` angular momentum matrix :math:`H(q)`, providing the linear
mapping from generalized velocities to subtree angular momentum :math:`h = H \dot q`. Contribution by
:github:user:`v-r-a`.
MJX
^^^
3. Improved performance of getting and putting device data.
4. Improved performance of getting and putting device data.
- Use ``tobytes()`` for numpy array serialization, which is orders of magnitude faster than converting to tuples.
- Avoid reallocating host ``mjData`` arrays when array shapes are unchanged.
- Speed up calculation of ``mjx.ncon`` for models with many geoms.
- Avoid calling ``mjx.ncon`` in ``mjx.get_data_into`` when ``nc`` can be derived from ``mjx.Data``.
4. Fixed a bug in ``mjx-viewer`` that prevented it from running. Updated ``mjx-viewer`` to use newer
5. Fixed a bug in ``mjx-viewer`` that prevented it from running. Updated ``mjx-viewer`` to use newer
``mjx.get_data_into`` function call.
5. Fixed a bug in ``mjx.euler`` that applied incorrect damping when using dense mass matrices.
6. Fixed a bug in ``mjx.solve`` that was causing slow convergence when using ``mjSOL_NEWTON`` in :ref:`mjtSolver`.
7. Added support for :ref:`mjOption.impratio<mjOption>` to ``mjx.Model``.
8. Added support for cameras in ``mjx.Model`` and ``mjx.Data``. Fixes :github:issue:`1422`.
9. Added an implementation of broadphase using `top_k` and bounding spheres.
6. Fixed a bug in ``mjx.euler`` that applied incorrect damping when using dense mass matrices.
7. Fixed a bug in ``mjx.solve`` that was causing slow convergence when using ``mjSOL_NEWTON`` in :ref:`mjtSolver`.
8. Added support for :ref:`mjOption.impratio<mjOption>` to ``mjx.Model``.
9. Added support for cameras in ``mjx.Model`` and ``mjx.Data``. Fixes :github:issue:`1422`.
10. Added an implementation of broadphase using `top_k` and bounding spheres.
Python bindings
^^^^^^^^^^^^^^^
10. Fixed incorrect data types in the bindings for the ``geom``, ``vert``, ``elem``, and ``flex`` array members
11. Fixed incorrect data types in the bindings for the ``geom``, ``vert``, ``elem``, and ``flex`` array members
of the ``mjContact`` struct, and all array members of the ``mjrContext`` struct.
doc/includes/references.h
+1
View File
@@ -2517,6 +2517,7 @@ void mj_jacGeom(const mjModel* m, const mjData* d, mjtNum* jacp, mjtNum* jacr, i
void mj_jacSite(const mjModel* m, const mjData* d, mjtNum* jacp, mjtNum* jacr, int site);
void mj_jacPointAxis(const mjModel* m, mjData* d, mjtNum* jacPoint, mjtNum* jacAxis,
const mjtNum point[3], const mjtNum axis[3], int body);
void mj_angmomMat(const mjModel* m, mjData* d, mjtNum* mat, int body);
int mj_name2id(const mjModel* m, int type, const char* name);
const char* mj_id2name(const mjModel* m, int type, int id);
void mj_fullM(const mjModel* m, mjtNum* dst, const mjtNum* M);
include/mujoco/mujoco.h
+3
View File
@@ -430,6 +430,9 @@ MJAPI void mj_jacSite(const mjModel* m, const mjData* d, mjtNum* jacp, mjtNum* j
MJAPI void mj_jacPointAxis(const mjModel* m, mjData* d, mjtNum* jacPoint, mjtNum* jacAxis,
const mjtNum point[3], const mjtNum axis[3], int body);
// Compute subtree angular momentum matrix.
MJAPI void mj_angmomMat(const mjModel* m, mjData* d, mjtNum* mat, int body);
// Get id of object with the specified mjtObj type and name, returns -1 if id not found.
MJAPI int mj_name2id(const mjModel* m, int type, const char* name);
introspect/functions.py
+30
View File
@@ -2402,6 +2402,36 @@ FUNCTIONS: Mapping[str, FunctionDecl] = dict([
),
doc='Compute translation end-effector Jacobian of point, and rotation Jacobian of axis.', # pylint: disable=line-too-long
)),
('mj_angmomMat',
FunctionDecl(
name='mj_angmomMat',
return_type=ValueType(name='void'),
parameters=(
FunctionParameterDecl(
name='m',
type=PointerType(
inner_type=ValueType(name='mjModel', is_const=True),
),
),
FunctionParameterDecl(
name='d',
type=PointerType(
inner_type=ValueType(name='mjData'),
),
),
FunctionParameterDecl(
name='mat',
type=PointerType(
inner_type=ValueType(name='mjtNum'),
),
),
FunctionParameterDecl(
name='body',
type=ValueType(name='int'),
),
),
doc='Compute subtree angular momentum matrix.',
)),
('mj_name2id',
FunctionDecl(
name='mj_name2id',
python/mujoco/bindings_test.py
+10
View File
@@ -744,6 +744,16 @@ class MuJoCoBindingsTest(parameterized.TestCase):
# Expect next states to be equal.
np.testing.assert_array_equal(state1a, state1b)
def test_mj_angmomMat(self): # pylint: disable=invalid-name
self.data.qvel = np.ones(self.model.nv, np.float64)
mujoco.mj_forward(self.model, self.data)
mujoco.mj_subtreeVel(self.model, self.data)
mat = np.empty((3, 10), np.float64)
mujoco.mj_angmomMat(self.model, self.data, mat, 0)
np.testing.assert_almost_equal(mat @ self.data.qvel,
self.data.subtree_angmom[0, :])
def test_mj_jacSite(self): # pylint: disable=invalid-name
mujoco.mj_forward(self.model, self.data)
site_id = mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_SITE, 'mysite')
python/mujoco/functions.cc
+8
View File
@@ -444,6 +444,14 @@ PYBIND11_MODULE(_functions, pymodule) {
jacr.has_value() ? jacr->data() : nullptr,
&(*point)[0], &(*axis)[0], body);
});
Def<traits::mj_angmomMat>(
pymodule, [](const raw::MjModel* m, raw::MjData* d,
Eigen::Ref<EigenArrayXX> mat, int body) {
if (mat.rows() != 3 || mat.cols() != m->nv) {
throw py::type_error("mat should be of shape (3, nv)");
}
return InterceptMjErrors(::mj_angmomMat)(m, d, mat.data(), body);
});
Def<traits::mj_name2id>(pymodule);
Def<traits::mj_id2name>(pymodule);
Def<traits::mj_fullM>(
src/engine/engine_support.c
+70
View File
@@ -776,6 +776,76 @@ int mj_jacSum(const mjModel* m, mjData* d, int* chain,
// compute subtree angular momentum matrix
void mj_angmomMat(const mjModel* m, mjData* d, mjtNum* mat, int body) {
int nv = m->nv;
mj_markStack(d);
// stack allocations
mjtNum* jacp = mj_stackAllocNum(d, 3*nv);
mjtNum* jacr = mj_stackAllocNum(d, 3*nv);
mjtNum* term1 = mj_stackAllocNum(d, 3*nv);
mjtNum* term2 = mj_stackAllocNum(d, 3*nv);
// clear output
mju_zero(mat, 3*nv);
// save the location of the subtree COM
mjtNum subtree_com[3];
mju_copy3(subtree_com, d->subtree_com+3*body);
for (int b=body; b < m->nbody; b++) {
// end of body subtree, break from the loop
if (b > body && m->body_parentid[b] < body) {
break;
}
// linear and angular velocity Jacobian of the body COM (inertial frame)
mj_jacBodyCom(m, d, jacp, jacr, b);
// orientation of the COM (inertial) frame of b-th body
mjtNum ximat[9];
mju_copy(ximat, d->ximat+9*b, 9);
// save the inertia matrix of b-th body
mjtNum inertia[9] = {0};
inertia[0] = m->body_inertia[3*b]; // inertia(1,1)
inertia[4] = m->body_inertia[3*b+1]; // inertia(2,2)
inertia[8] = m->body_inertia[3*b+2]; // inertia(3,3)
// term1 = body angular momentum about self COM in world frame
mjtNum tmp1[9], tmp2[9];
mju_mulMatMat(tmp1, ximat, inertia, 3, 3, 3); // tmp1 = ximat * inertia
mju_mulMatMatT(tmp2, tmp1, ximat, 3, 3, 3); // tmp2 = ximat * inertia * ximat^T
mju_mulMatMat(term1, tmp2, jacr, 3, 3, nv); // term1 = ximat * inertia * ximat^T * jacr
// location of body COM w.r.t subtree COM
mjtNum com[3];
mju_sub3(com, d->xipos+3*b, subtree_com);
// skew symmetric matrix representing body_com vector
mjtNum com_mat[9] = {0};
com_mat[1] = -com[2];
com_mat[2] = com[1];
com_mat[3] = com[2];
com_mat[5] = -com[0];
com_mat[6] = -com[1];
com_mat[7] = com[0];
// term2 = moment of linear momentum
mju_mulMatMat(term2, com_mat, jacp, 3, 3, nv); // term2 = com_mat * jacp
mju_scl(term2, term2, m->body_mass[b], 3 * nv); // term2 = com_mat * jacp * mass
// mat += term1 + term2
mju_addTo(mat, term1, 3*nv);
mju_addTo(mat, term2, 3*nv);
}
mj_freeStack(d);
}
//-------------------------- name functions --------------------------------------------------------
// get number of objects and name addresses for given object type
src/engine/engine_support.h
+3
View File
@@ -107,6 +107,9 @@ int mj_jacSum(const mjModel* m, mjData* d, int* chain,
int n, const int* body, const mjtNum* weight,
const mjtNum point[3], mjtNum* jac, int flg_rot);
// compute subtree angular momentum matrix
MJAPI void mj_angmomMat(const mjModel* m, mjData* d, mjtNum* mat, int body);
//-------------------------- name functions --------------------------------------------------------
test/engine/engine_support_test.cc
+121 -1
View File
@@ -16,6 +16,7 @@
#include "src/engine/engine_support.h"
#include <limits>
#include <random>
#include <string>
#include <vector>
@@ -37,6 +38,125 @@ using ::testing::ContainsRegex;
using ::testing::MatchesRegex;
using ::testing::Pointwise;
using ::testing::ElementsAreArray;
using AngMomMatTest = MujocoTest;
static constexpr char AngMomTestingModel[] = R"(
<mujoco>
<option>
<flag gravity="disable"/>
</option>
<worldbody>
<body name="link1" pos="0 0 0.5">
<freejoint/>
<geom type="ellipsoid" size="0.15 0.17 0.19" quat="1 .2 .3 .4"/>
<body name="link2" >
<joint type="hinge" axis="1 0 0" />
<geom type="capsule" size="0.05" fromto="0 0 0 0 0.5 0"/>
<body pos="0 0.6 0">
<joint type="slide" axis="1 0 0"/>
<geom type="capsule" size="0.05 0.2" quat="0.707 0 0.707 0"/>
<body name="link3">
<joint type="ball" pos="0.2 0 0"/>
<geom type="capsule" pos="0.2 0 0" size="0.03 0.4"/>
</body>
</body>
</body>
</body>
</worldbody>
<keyframe>
<key qvel="0 0 0 .1 .2 .3 .4 .5 .4 .3 .2"/>
</keyframe>
</mujoco>
)";
// compare subtree angular momentum computed in two ways
TEST_F(AngMomMatTest, CompareAngMom) {
mjModel* model = LoadModelFromString(AngMomTestingModel);
int nv = model->nv;
int bodyid = mj_name2id(model, mjOBJ_BODY, "link1");
mjData* data = mj_makeData(model);
// reset to the keyframe with some angular velocities
mj_resetDataKeyframe(model, data, 0);
mj_forward(model, data);
// get the reference value of angular momentum
mj_subtreeVel(model, data);
mjtNum angmom_ref[3];
mju_copy3(angmom_ref, data->subtree_angmom+3*bodyid);
// compute angular momentum using the angular momentum matrix
mjtNum* angmom_mat = (mjtNum*) mju_malloc(sizeof(mjtNum)*3*nv);
mj_angmomMat(model, data, angmom_mat, bodyid);
mjtNum angmom_test[3];
mju_mulMatVec(angmom_test, angmom_mat, data->qvel, 3, nv);
// compare the two angular momentum values
static const mjtNum tol = 1e-8;
for (int i = 0; i < 3; i++) {
EXPECT_THAT(angmom_ref[i], DoubleNear(angmom_test[i], tol));
}
mju_free(angmom_mat);
mj_deleteData(data);
mj_deleteModel(model);
}
// compare subtree angular momentum matrix: analytical and findiff
TEST_F(AngMomMatTest, CompareAngMomMats) {
mjModel* model = LoadModelFromString(AngMomTestingModel);
int nv = model->nv;
int bodyid = mj_name2id(model, mjOBJ_BODY, "link1");
mjData* data = mj_makeData(model);
mjtNum* angmom_mat = (mjtNum*) mju_malloc(sizeof(mjtNum)*3*nv);
mjtNum* angmom_mat_fd = (mjtNum*) mju_malloc(sizeof(mjtNum)*3*nv);
// reset to the keyframe with some angular velocities
mj_resetDataKeyframe(model, data, 0);
mj_forward(model, data);
// compute the angular momentum matrix using the analytical method
mj_angmomMat(model, data, angmom_mat, bodyid);
// compute the angular momentum matrix using finite differences
static const mjtNum eps = 1e-6;
for (int i = 0; i < nv; i++) {
// reset vel, forward nudge i-th dof, get angmom
mju_copy(data->qvel, model->key_qvel, model->nv);
data->qvel[i] += eps;
mj_forward(model, data);
mj_subtreeVel(model, data);
mjtNum agmf[3];
mju_copy3(agmf, data->subtree_angmom+3*bodyid);
// reset vel, backward nudge i-th dof, get angmom
mju_copy(data->qvel, model->key_qvel, model->nv);
data->qvel[i] -= eps;
mj_forward(model, data);
mj_subtreeVel(model, data);
mjtNum agmb[3];
mju_copy3(agmb, data->subtree_angmom+3*bodyid);
// finite-difference the angmom matrix
for (int j = 0; j < 3; j++) {
angmom_mat_fd[nv*j+i] = (agmf[j] - agmb[j]) / (2 * eps);
}
}
// compare the two matrices
static const mjtNum tol = 1e-8;
for (int i = 0; i < 3*nv; i++) {
EXPECT_THAT(angmom_mat_fd[i], DoubleNear(angmom_mat[i], tol));
}
mju_free(angmom_mat_fd);
mju_free(angmom_mat);
mj_deleteData(data);
mj_deleteModel(model);
}
using JacobianTest = MujocoTest;
static const mjtNum max_abs_err = std::numeric_limits<float>::epsilon();
@@ -446,7 +566,7 @@ TEST_F(AddMTest, DenseSameAsSparse) {
rowadr.data(), colind.data());
// dense addM
mj_addM(m, d, dst_dense.data(), NULL, NULL, NULL);
mj_addM(m, d, dst_dense.data(), nullptr, nullptr, nullptr);
// dense comparison, should be same matrix
EXPECT_THAT(dst_dense, ElementsAreArray(dst_sparse));
unity/Runtime/Bindings/MjBindings.cs
+3
View File
@@ -6627,6 +6627,9 @@ public static unsafe extern void mj_jacSite(mjModel_* m, mjData_* d, double* jac
[DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)]
public static unsafe extern void mj_jacPointAxis(mjModel_* m, mjData_* d, double* jacPoint, double* jacAxis, double* point, double* axis, int body);
[DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)]
public static unsafe extern void mj_angmomMat(mjModel_* m, mjData_* d, double* mat, int body);
[DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)]
public static unsafe extern int mj_name2id(mjModel_* m, int type, [MarshalAs(UnmanagedType.LPStr)]string name);