Extract memory allocation functions and core utilities

PiperOrigin-RevId: 801745499
Change-Id: Iaf05c3430769d3115743d8ab020d13148cb2eb59
This commit is contained in:
Yuval Tassa
2025-09-01 03:37:16 -07:00
committed by Copybara-Service
parent 5d598a49a2
commit b9900db00e
35 changed files with 1652 additions and 1445 deletions
+4
View File
@@ -29,6 +29,8 @@ set(MUJOCO_ENGINE_SRCS
engine_collision_sdf.h
engine_core_constraint.c
engine_core_constraint.h
engine_core_util.c
engine_core_util.h
engine_core_smooth.c
engine_core_smooth.h
engine_crossplatform.cc
@@ -47,6 +49,8 @@ set(MUJOCO_ENGINE_SRCS
engine_io.c
engine_io.h
engine_macro.h
engine_memory.c
engine_memory.h
engine_name.c
engine_name.h
engine_passive.c
+1 -1
View File
@@ -25,7 +25,7 @@
#include <mujoco/mjmodel.h>
#include "engine/engine_collision_gjk.h"
#include "engine/engine_collision_primitive.h"
#include "engine/engine_io.h"
#include "engine/engine_memory.h"
#include "engine/engine_util_blas.h"
#include "engine/engine_util_errmem.h"
#include "engine/engine_util_misc.h"
+2 -1
View File
@@ -26,10 +26,11 @@
#include "engine/engine_collision_primitive.h"
#include "engine/engine_collision_sdf.h"
#include "engine/engine_core_constraint.h"
#include "engine/engine_core_util.h"
#include "engine/engine_io.h"
#include "engine/engine_macro.h"
#include "engine/engine_memory.h"
#include "engine/engine_sort.h"
#include "engine/engine_support.h"
#include "engine/engine_util_blas.h"
#include "engine/engine_util_errmem.h"
#include "engine/engine_util_misc.h"
+3 -25
View File
@@ -22,9 +22,9 @@
#include <mujoco/mjmodel.h>
#include <mujoco/mjsan.h> // IWYU pragma: keep
#include <mujoco/mjxmacro.h>
#include "engine/engine_core_util.h"
#include "engine/engine_core_smooth.h"
#include "engine/engine_io.h"
#include "engine/engine_support.h"
#include "engine/engine_memory.h"
#include "engine/engine_util_blas.h"
#include "engine/engine_util_errmem.h"
#include "engine/engine_util_misc.h"
@@ -84,29 +84,6 @@ static int arenaAllocEfc(const mjModel* m, mjData* d) {
// determine type of friction cone
int mj_isPyramidal(const mjModel* m) {
if (m->opt.cone == mjCONE_PYRAMIDAL) {
return 1;
} else {
return 0;
}
}
// determine type of constraint Jacobian
int mj_isSparse(const mjModel* m) {
if (m->opt.jacobian == mjJAC_SPARSE ||
(m->opt.jacobian == mjJAC_AUTO && m->nv >= 60)) {
return 1;
} else {
return 0;
}
}
// determine type of solver
int mj_isDual(const mjModel* m) {
if (m->opt.solver == mjSOL_PGS || m->opt.noslip_iterations > 0) {
@@ -133,6 +110,7 @@ void mj_assignFriction(const mjModel* m, mjtNum* target, const mjtNum* source) {
// assign/override contact reference parameters
void mj_assignRef(const mjModel* m, mjtNum* target, const mjtNum* source) {
if (mjENABLED(mjENBL_OVERRIDE)) {
-7
View File
@@ -27,12 +27,6 @@ extern "C" {
//-------------------------- Jacobian-related ------------------------------------------------------
// determine type of friction cone
MJAPI int mj_isPyramidal(const mjModel* m);
// determine type of constraint Jacobian
MJAPI int mj_isSparse(const mjModel* m);
// determine type of solver
MJAPI int mj_isDual(const mjModel* m);
@@ -60,7 +54,6 @@ mjtNum mj_assignMargin(const mjModel* m, mjtNum source);
// add contact to d->contact list; return 0 if success; 1 if buffer full
MJAPI int mj_addContact(const mjModel* m, mjData* d, const mjContact* con);
//-------------------------- constraint instantiation ----------------------------------------------
// equality constraints
+2 -2
View File
@@ -22,10 +22,10 @@
#include <mujoco/mjmodel.h>
#include <mujoco/mjsan.h> // IWYU pragma: keep
#include "engine/engine_core_constraint.h"
#include "engine/engine_core_util.h"
#include "engine/engine_crossplatform.h"
#include "engine/engine_io.h"
#include "engine/engine_macro.h"
#include "engine/engine_support.h"
#include "engine/engine_memory.h"
#include "engine/engine_util_blas.h"
#include "engine/engine_util_errmem.h"
#include "engine/engine_util_misc.h"
+961
View File
@@ -0,0 +1,961 @@
// Copyright 2025 DeepMind Technologies Limited
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#include "engine/engine_core_util.h"
#include <stddef.h>
#include <mujoco/mjdata.h>
#include <mujoco/mjmodel.h>
#include "engine/engine_crossplatform.h"
#include "engine/engine_memory.h"
#include "engine/engine_util_blas.h"
#include "engine/engine_util_errmem.h"
#include "engine/engine_util_misc.h"
#include "engine/engine_util_sparse.h"
#include "engine/engine_util_spatial.h"
// determine type of constraint Jacobian
int mj_isSparse(const mjModel* m) {
if (m->opt.jacobian == mjJAC_SPARSE ||
(m->opt.jacobian == mjJAC_AUTO && m->nv >= 60)) {
return 1;
} else {
return 0;
}
}
// determine type of friction cone
int mj_isPyramidal(const mjModel* m) {
if (m->opt.cone == mjCONE_PYRAMIDAL) {
return 1;
} else {
return 0;
}
}
//-------------------------- sparse chains ---------------------------------------------------------
// merge dof chains for two bodies
int mj_mergeChain(const mjModel* m, int* chain, int b1, int b2) {
int da1, da2, NV = 0;
// skip fixed bodies
while (b1 && !m->body_dofnum[b1]) {
b1 = m->body_parentid[b1];
}
while (b2 && !m->body_dofnum[b2]) {
b2 = m->body_parentid[b2];
}
// neither body is movable: empty chain
if (b1 == 0 && b2 == 0) {
return 0;
}
// initialize last dof address for each body
da1 = m->body_dofadr[b1] + m->body_dofnum[b1] - 1;
da2 = m->body_dofadr[b2] + m->body_dofnum[b2] - 1;
// merge chains
while (da1 >= 0 || da2 >= 0) {
chain[NV] = mjMAX(da1, da2);
if (da1 == chain[NV]) {
da1 = m->dof_parentid[da1];
}
if (da2 == chain[NV]) {
da2 = m->dof_parentid[da2];
}
NV++;
}
// reverse order of chain: make it increasing
for (int i=0; i < NV/2; i++) {
int tmp = chain[i];
chain[i] = chain[NV-i-1];
chain[NV-i-1] = tmp;
}
return NV;
}
// merge dof chains for two simple bodies
int mj_mergeChainSimple(const mjModel* m, int* chain, int b1, int b2) {
// swap bodies if wrong order
if (b1 > b2) {
int tmp = b1;
b1 = b2;
b2 = tmp;
}
// init
int n1 = m->body_dofnum[b1], n2 = m->body_dofnum[b2];
// both fixed: nothing to do
if (n1 == 0 && n2 == 0) {
return 0;
}
// copy b1 dofs
for (int i=0; i < n1; i++) {
chain[i] = m->body_dofadr[b1] + i;
}
// copy b2 dofs
for (int i=0; i < n2; i++) {
chain[n1+i] = m->body_dofadr[b2] + i;
}
return (n1+n2);
}
// get body chain
int mj_bodyChain(const mjModel* m, int body, int* chain) {
// simple body
if (m->body_simple[body]) {
int dofnum = m->body_dofnum[body];
for (int i=0; i < dofnum; i++) {
chain[i] = m->body_dofadr[body] + i;
}
return dofnum;
}
// general case
else {
// skip fixed bodies
while (body && !m->body_dofnum[body]) {
body = m->body_parentid[body];
}
// not movable: empty chain
if (body == 0) {
return 0;
}
// initialize last dof
int da = m->body_dofadr[body] + m->body_dofnum[body] - 1;
int NV = 0;
// construct chain from child to parent
while (da >= 0) {
chain[NV++] = da;
da = m->dof_parentid[da];
}
// reverse order of chain: make it increasing
for (int i=0; i < NV/2; i++) {
int tmp = chain[i];
chain[i] = chain[NV-i-1];
chain[NV-i-1] = tmp;
}
return NV;
}
}
//-------------------------- Jacobians -------------------------------------------------------------
// compute 3/6-by-nv Jacobian of global point attached to given body
void mj_jac(const mjModel* m, const mjData* d,
mjtNum* jacp, mjtNum* jacr, const mjtNum point[3], int body) {
int nv = m->nv;
mjtNum offset[3];
// clear jacobians, compute offset if required
if (jacp) {
mju_zero(jacp, 3*nv);
mju_sub3(offset, point, d->subtree_com+3*m->body_rootid[body]);
}
if (jacr) {
mju_zero(jacr, 3*nv);
}
// skip fixed bodies
while (body && !m->body_dofnum[body]) {
body = m->body_parentid[body];
}
// no movable body found: nothing to do
if (!body) {
return;
}
// get last dof that affects this (as well as the original) body
int i = m->body_dofadr[body] + m->body_dofnum[body] - 1;
// backward pass over dof ancestor chain
while (i >= 0) {
mjtNum* cdof = d->cdof+6*i;
// construct rotation jacobian
if (jacr) {
jacr[i+0*nv] = cdof[0];
jacr[i+1*nv] = cdof[1];
jacr[i+2*nv] = cdof[2];
}
// construct translation jacobian (correct for rotation)
if (jacp) {
mjtNum tmp[3];
mju_cross(tmp, cdof, offset);
jacp[i+0*nv] = cdof[3] + tmp[0];
jacp[i+1*nv] = cdof[4] + tmp[1];
jacp[i+2*nv] = cdof[5] + tmp[2];
}
// advance to parent dof
i = m->dof_parentid[i];
}
}
// compute body Jacobian
void mj_jacBody(const mjModel* m, const mjData* d, mjtNum* jacp, mjtNum* jacr, int body) {
mj_jac(m, d, jacp, jacr, d->xpos+3*body, body);
}
// compute body-com Jacobian
void mj_jacBodyCom(const mjModel* m, const mjData* d, mjtNum* jacp, mjtNum* jacr, int body) {
mj_jac(m, d, jacp, jacr, d->xipos+3*body, body);
}
// compute subtree-com Jacobian
void mj_jacSubtreeCom(const mjModel* m, mjData* d, mjtNum* jacp, int body) {
int nv = m->nv;
mj_markStack(d);
mjtNum* jacp_b = mjSTACKALLOC(d, 3*nv, mjtNum);
// clear output
mju_zero(jacp, 3*nv);
// forward pass starting from 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;
}
// b is in the body subtree, add mass-weighted Jacobian into jacp
mj_jac(m, d, jacp_b, NULL, d->xipos+3*b, b);
mju_addToScl(jacp, jacp_b, m->body_mass[b], 3*nv);
}
// normalize by subtree mass
mju_scl(jacp, jacp, 1/m->body_subtreemass[body], 3*nv);
mj_freeStack(d);
}
// compute geom Jacobian
void mj_jacGeom(const mjModel* m, const mjData* d, mjtNum* jacp, mjtNum* jacr, int geom) {
mj_jac(m, d, jacp, jacr, d->geom_xpos + 3*geom, m->geom_bodyid[geom]);
}
// compute site Jacobian
void mj_jacSite(const mjModel* m, const mjData* d, mjtNum* jacp, mjtNum* jacr, int site) {
mj_jac(m, d, jacp, jacr, d->site_xpos + 3*site, m->site_bodyid[site]);
}
// compute translation Jacobian of point, and rotation Jacobian of axis
void mj_jacPointAxis(const mjModel* m, mjData* d, mjtNum* jacPoint, mjtNum* jacAxis,
const mjtNum point[3], const mjtNum axis[3], int body) {
int nv = m->nv;
// get full Jacobian of point
mj_markStack(d);
mjtNum* jacp = (jacPoint ? jacPoint : mjSTACKALLOC(d, 3*nv, mjtNum));
mjtNum* jacr = mjSTACKALLOC(d, 3*nv, mjtNum);
mj_jac(m, d, jacp, jacr, point, body);
// jacAxis_col = cross(jacr_col, axis)
if (jacAxis) {
for (int i=0; i < nv; i++) {
jacAxis[ i] = jacr[ nv+i]*axis[2] - jacr[2*nv+i]*axis[1];
jacAxis[ nv+i] = jacr[2*nv+i]*axis[0] - jacr[ i]*axis[2];
jacAxis[2*nv+i] = jacr[ i]*axis[1] - jacr[ nv+i]*axis[0];
}
}
mj_freeStack(d);
}
// compute 3/6-by-nv sparse Jacobian of global point attached to given body
void mj_jacSparse(const mjModel* m, const mjData* d,
mjtNum* jacp, mjtNum* jacr, const mjtNum* point, int body,
int NV, const int* chain) {
// clear jacobians
if (jacp) {
mju_zero(jacp, 3*NV);
}
if (jacr) {
mju_zero(jacr, 3*NV);
}
// compute point-com offset
mjtNum offset[3];
mju_sub3(offset, point, d->subtree_com+3*m->body_rootid[body]);
// skip fixed bodies
while (body && !m->body_dofnum[body]) {
body = m->body_parentid[body];
}
// no movable body found: nothing to do
if (!body) {
return;
}
// get last dof that affects this (as well as the original) body
int da = m->body_dofadr[body] + m->body_dofnum[body] - 1;
// start and the end of the chain (chain is in increasing order)
int ci = NV-1;
// backward pass over dof ancestor chain
while (da >= 0) {
// find chain index for this dof
while (ci >= 0 && chain[ci] > da) {
ci--;
}
// make sure we found it; SHOULD NOT OCCUR
if (ci < 0 || chain[ci] != da) {
mjERROR("dof index %d not found in chain", da);
}
const mjtNum* cdof = d->cdof + 6*da;
// construct rotation jacobian
if (jacr) {
jacr[ci+0*NV] = cdof[0];
jacr[ci+1*NV] = cdof[1];
jacr[ci+2*NV] = cdof[2];
}
// construct translation jacobian (correct for rotation)
if (jacp) {
mjtNum tmp[3];
mju_cross(tmp, cdof, offset);
jacp[ci+0*NV] = cdof[3] + tmp[0];
jacp[ci+1*NV] = cdof[4] + tmp[1];
jacp[ci+2*NV] = cdof[5] + tmp[2];
}
// advance to parent dof
da = m->dof_parentid[da];
}
}
// sparse Jacobian difference for simple body contacts
void mj_jacSparseSimple(const mjModel* m, const mjData* d,
mjtNum* jacdifp, mjtNum* jacdifr, const mjtNum* point,
int body, int flg_second, int NV, int start) {
// compute point-com offset
mjtNum offset[3];
mju_sub3(offset, point, d->subtree_com+3*m->body_rootid[body]);
// skip fixed body
if (!m->body_dofnum[body]) {
return;
}
// process dofs
int ci = start;
int end = m->body_dofadr[body] + m->body_dofnum[body];
for (int da=m->body_dofadr[body]; da < end; da++) {
mjtNum *cdof = d->cdof+6*da;
// construct rotation jacobian
if (jacdifr) {
// plus sign
if (flg_second) {
jacdifr[ci+0*NV] = cdof[0];
jacdifr[ci+1*NV] = cdof[1];
jacdifr[ci+2*NV] = cdof[2];
}
// minus sign
else {
jacdifr[ci+0*NV] = -cdof[0];
jacdifr[ci+1*NV] = -cdof[1];
jacdifr[ci+2*NV] = -cdof[2];
}
}
// construct translation jacobian (correct for rotation)
if (jacdifp) {
mjtNum tmp[3];
mju_cross(tmp, cdof, offset);
// plus sign
if (flg_second) {
jacdifp[ci+0*NV] = (cdof[3] + tmp[0]);
jacdifp[ci+1*NV] = (cdof[4] + tmp[1]);
jacdifp[ci+2*NV] = (cdof[5] + tmp[2]);
}
// minus sign
else {
jacdifp[ci+0*NV] = -(cdof[3] + tmp[0]);
jacdifp[ci+1*NV] = -(cdof[4] + tmp[1]);
jacdifp[ci+2*NV] = -(cdof[5] + tmp[2]);
}
}
// advance jacdif counter
ci++;
}
}
// dense or sparse Jacobian difference for two body points: pos2 - pos1, global
int mj_jacDifPair(const mjModel* m, const mjData* d, int* chain,
int b1, int b2, const mjtNum pos1[3], const mjtNum pos2[3],
mjtNum* jac1p, mjtNum* jac2p, mjtNum* jacdifp,
mjtNum* jac1r, mjtNum* jac2r, mjtNum* jacdifr) {
int issimple = (m->body_simple[b1] && m->body_simple[b2]);
int issparse = mj_isSparse(m);
int NV = m->nv;
// skip if no DOFs
if (!NV) {
return 0;
}
// construct merged chain of body dofs
if (issparse) {
if (issimple) {
NV = mj_mergeChainSimple(m, chain, b1, b2);
} else {
NV = mj_mergeChain(m, chain, b1, b2);
}
}
// skip if empty chain
if (!NV) {
return 0;
}
// sparse case
if (issparse) {
// simple: fast processing
if (issimple) {
// first body
mj_jacSparseSimple(m, d, jacdifp, jacdifr, pos1, b1, 0, NV,
b1 < b2 ? 0 : m->body_dofnum[b2]);
// second body
mj_jacSparseSimple(m, d, jacdifp, jacdifr, pos2, b2, 1, NV,
b2 < b1 ? 0 : m->body_dofnum[b1]);
}
// regular processing
else {
// Jacobians
mj_jacSparse(m, d, jac1p, jac1r, pos1, b1, NV, chain);
mj_jacSparse(m, d, jac2p, jac2r, pos2, b2, NV, chain);
// differences
if (jacdifp) {
mju_sub(jacdifp, jac2p, jac1p, 3*NV);
}
if (jacdifr) {
mju_sub(jacdifr, jac2r, jac1r, 3*NV);
}
}
}
// dense case
else {
// Jacobians
mj_jac(m, d, jac1p, jac1r, pos1, b1);
mj_jac(m, d, jac2p, jac2r, pos2, b2);
// differences
if (jacdifp) {
mju_sub(jacdifp, jac2p, jac1p, 3*NV);
}
if (jacdifr) {
mju_sub(jacdifr, jac2r, jac1r, 3*NV);
}
}
return NV;
}
// dense or sparse weighted sum of multiple body Jacobians at same point
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) {
int nv = m->nv, NV;
mjtNum* jacp = jac;
mjtNum* jacr = flg_rot ? jac + 3*nv : NULL;
mj_markStack(d);
mjtNum* jtmp = mjSTACKALLOC(d, flg_rot ? 6*nv : 3*nv, mjtNum);
mjtNum* jp = jtmp;
mjtNum* jr = flg_rot ? jtmp + 3*nv : NULL;
// sparse
if (mj_isSparse(m)) {
mjtNum* buf = mjSTACKALLOC(d, flg_rot ? 6*nv : 3*nv, mjtNum);
int* buf_ind = mjSTACKALLOC(d, nv, int);
int* bodychain = mjSTACKALLOC(d, nv, int);
// set first
NV = mj_bodyChain(m, body[0], chain);
if (NV) {
// get Jacobian
if (m->body_simple[body[0]]) {
mj_jacSparseSimple(m, d, jacp, jacr, point, body[0], 1, NV, 0);
} else {
mj_jacSparse(m, d, jacp, jacr, point, body[0], NV, chain);
}
// apply weight
mju_scl(jac, jac, weight[0], flg_rot ? 6*NV : 3*NV);
}
// accumulate remaining
for (int i=1; i < n; i++) {
// get body chain and Jacobian
int bodyNV = mj_bodyChain(m, body[i], bodychain);
if (!bodyNV) {
continue;
}
if (m->body_simple[body[i]]) {
mj_jacSparseSimple(m, d, jp, jr, point, body[i], 1, bodyNV, 0);
} else {
mj_jacSparse(m, d, jp, jr, point, body[i], bodyNV, bodychain);
}
// combine sparse matrices
NV = mju_addToSparseMat(jac, jtmp, nv, flg_rot ? 6 : 3, weight[i],
NV, bodyNV, chain, bodychain, buf, buf_ind);
}
}
// dense
else {
// set first
mj_jac(m, d, jacp, jacr, point, body[0]);
mju_scl(jac, jac, weight[0], flg_rot ? 6*nv : 3*nv);
// accumulate remaining
for (int i=1; i < n; i++) {
mj_jac(m, d, jp, jr, point, body[i]);
mju_addToScl(jac, jtmp, weight[i], flg_rot ? 6*nv : 3*nv);
}
NV = nv;
}
mj_freeStack(d);
return NV;
}
// compute 3/6-by-nv Jacobian time derivative of global point attached to given body
void mj_jacDot(const mjModel* m, const mjData* d,
mjtNum* jacp, mjtNum* jacr, const mjtNum point[3], int body) {
int nv = m->nv;
mjtNum offset[3];
mjtNum pvel[6]; // point velocity (rot:lin order)
// clear jacobians, compute offset and pvel if required
if (jacp) {
mju_zero(jacp, 3*nv);
const mjtNum* com = d->subtree_com+3*m->body_rootid[body];
mju_sub3(offset, point, com);
mju_transformSpatial(pvel, d->cvel+6*body, 0, point, com, 0);
}
if (jacr) {
mju_zero(jacr, 3*nv);
}
// skip fixed bodies
while (body && !m->body_dofnum[body]) {
body = m->body_parentid[body];
}
// no movable body found: nothing to do
if (!body) {
return;
}
// get last dof that affects this (as well as the original) body
int i = m->body_dofadr[body] + m->body_dofnum[body] - 1;
// backward pass over dof ancestor chain
while (i >= 0) {
mjtNum cdof_dot[6];
mju_copy(cdof_dot, d->cdof_dot+6*i, 6);
mjtNum* cdof = d->cdof+6*i;
// check for quaternion
mjtJoint type = m->jnt_type[m->dof_jntid[i]];
int dofadr = m->jnt_dofadr[m->dof_jntid[i]];
int is_quat = type == mjJNT_BALL || (type == mjJNT_FREE && i >= dofadr + 3);
// compute cdof_dot for quaternion (use current body cvel)
if (is_quat) {
mju_crossMotion(cdof_dot, d->cvel+6*m->dof_bodyid[i], cdof);
}
// construct rotation jacobian
if (jacr) {
jacr[i+0*nv] += cdof_dot[0];
jacr[i+1*nv] += cdof_dot[1];
jacr[i+2*nv] += cdof_dot[2];
}
// construct translation jacobian (correct for rotation)
if (jacp) {
// first correction term, account for varying cdof
mjtNum tmp1[3];
mju_cross(tmp1, cdof_dot, offset);
// second correction term, account for point translational velocity
mjtNum tmp2[3];
mju_cross(tmp2, cdof, pvel + 3);
jacp[i+0*nv] += cdof_dot[3] + tmp1[0] + tmp2[0];
jacp[i+1*nv] += cdof_dot[4] + tmp1[1] + tmp2[1];
jacp[i+2*nv] += cdof_dot[5] + tmp1[2] + tmp2[2];
}
// advance to parent dof
i = m->dof_parentid[i];
}
}
// 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 = mjSTACKALLOC(d, 3*nv, mjtNum);
mjtNum* jacr = mjSTACKALLOC(d, 3*nv, mjtNum);
mjtNum* term1 = mjSTACKALLOC(d, 3*nv, mjtNum);
mjtNum* term2 = mjSTACKALLOC(d, 3*nv, mjtNum);
// 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_mulMatMat3(tmp1, ximat, inertia); // tmp1 = ximat * inertia
mju_mulMatMatT3(tmp2, tmp1, ximat); // 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);
}
// count warnings, print only the first time
void mj_warning(mjData* d, int warning, int info) {
// check type
if (warning < 0 || warning >= mjNWARNING) {
mjERROR("invalid warning type %d", warning);
}
// save info (override previous)
d->warning[warning].lastinfo = info;
// print message only the first time this warning is encountered
if (!d->warning[warning].number) {
mju_warning("%s Time = %.4f.", mju_warningText(warning, info), d->time);
}
// increase counter
d->warning[warning].number++;
}
// compute object 6D velocity in object-centered frame, world/local orientation
void mj_objectVelocity(const mjModel* m, const mjData* d,
int objtype, int objid, mjtNum res[6], int flg_local) {
int bodyid = 0;
const mjtNum *pos = 0, *rot = 0;
// body-inertial
if (objtype == mjOBJ_BODY) {
bodyid = objid;
pos = d->xipos+3*objid;
rot = (flg_local ? d->ximat+9*objid : 0);
}
// body-regular
else if (objtype == mjOBJ_XBODY) {
bodyid = objid;
pos = d->xpos+3*objid;
rot = (flg_local ? d->xmat+9*objid : 0);
}
// geom
else if (objtype == mjOBJ_GEOM) {
bodyid = m->geom_bodyid[objid];
pos = d->geom_xpos+3*objid;
rot = (flg_local ? d->geom_xmat+9*objid : 0);
}
// site
else if (objtype == mjOBJ_SITE) {
bodyid = m->site_bodyid[objid];
pos = d->site_xpos+3*objid;
rot = (flg_local ? d->site_xmat+9*objid : 0);
}
// camera
else if (objtype == mjOBJ_CAMERA) {
bodyid = m->cam_bodyid[objid];
pos = d->cam_xpos+3*objid;
rot = (flg_local ? d->cam_xmat+9*objid : 0);
}
// object without spatial frame
else {
mjERROR("invalid object type %d", objtype);
}
// transform velocity
mju_transformSpatial(res, d->cvel+6*bodyid, 0, pos, d->subtree_com+3*m->body_rootid[bodyid], rot);
}
// compute object 6D acceleration in object-centered frame, world/local orientation
void mj_objectAcceleration(const mjModel* m, const mjData* d,
int objtype, int objid, mjtNum res[6], int flg_local) {
int bodyid = 0;
const mjtNum *pos = 0, *rot = 0;
mjtNum correction[3], vel[6];
// body-inertial
if (objtype == mjOBJ_BODY) {
bodyid = objid;
pos = d->xipos+3*objid;
rot = (flg_local ? d->ximat+9*objid : 0);
}
// body-regular
else if (objtype == mjOBJ_XBODY) {
bodyid = objid;
pos = d->xpos+3*objid;
rot = (flg_local ? d->xmat+9*objid : 0);
}
// geom
else if (objtype == mjOBJ_GEOM) {
bodyid = m->geom_bodyid[objid];
pos = d->geom_xpos+3*objid;
rot = (flg_local ? d->geom_xmat+9*objid : 0);
}
// site
else if (objtype == mjOBJ_SITE) {
bodyid = m->site_bodyid[objid];
pos = d->site_xpos+3*objid;
rot = (flg_local ? d->site_xmat+9*objid : 0);
}
// camera
else if (objtype == mjOBJ_CAMERA) {
bodyid = m->cam_bodyid[objid];
pos = d->cam_xpos+3*objid;
rot = (flg_local ? d->cam_xmat+9*objid : 0);
}
// object without spatial frame
else {
mjERROR("invalid object type %d", objtype);
}
// transform com-based velocity to local frame
mju_transformSpatial(vel, d->cvel+6*bodyid, 0, pos, d->subtree_com+3*m->body_rootid[bodyid], rot);
// transform com-based acceleration to local frame
mju_transformSpatial(res, d->cacc+6*bodyid, 0, pos, d->subtree_com+3*m->body_rootid[bodyid], rot);
// acc_tran += vel_rot x vel_tran
mju_cross(correction, vel, vel+3);
mju_addTo3(res+3, correction);
}
// extract 6D force:torque for one contact, in contact frame
void mj_contactForce(const mjModel* m, const mjData* d, int id, mjtNum result[6]) {
mjContact* con;
// clear result
mju_zero(result, 6);
// make sure contact is valid
if (id >= 0 && id < d->ncon && d->contact[id].efc_address >= 0) {
// get contact pointer
con = d->contact + id;
if (mj_isPyramidal(m)) {
mju_decodePyramid(result, d->efc_force + con->efc_address, con->friction, con->dim);
} else {
mju_copy(result, d->efc_force + con->efc_address, con->dim);
}
}
}
// map from body local to global Cartesian coordinates
void mj_local2Global(mjData* d, mjtNum xpos[3], mjtNum xmat[9],
const mjtNum pos[3], const mjtNum quat[4],
int body, mjtByte sameframe) {
mjtSameFrame sf = sameframe;
// position
if (xpos && pos) {
switch (sf) {
case mjSAMEFRAME_NONE:
case mjSAMEFRAME_BODYROT:
case mjSAMEFRAME_INERTIAROT:
mju_mulMatVec3(xpos, d->xmat+9*body, pos);
mju_addTo3(xpos, d->xpos+3*body);
break;
case mjSAMEFRAME_BODY:
mju_copy3(xpos, d->xpos+3*body);
break;
case mjSAMEFRAME_INERTIA:
mju_copy3(xpos, d->xipos+3*body);
break;
}
}
// orientation
if (xmat && quat) {
mjtNum tmp[4];
switch (sf) {
case mjSAMEFRAME_NONE:
mju_mulQuat(tmp, d->xquat+4*body, quat);
mju_quat2Mat(xmat, tmp);
break;
case mjSAMEFRAME_BODY:
case mjSAMEFRAME_BODYROT:
mju_copy(xmat, d->xmat+9*body, 9);
break;
case mjSAMEFRAME_INERTIA:
case mjSAMEFRAME_INERTIAROT:
mju_copy(xmat, d->ximat+9*body, 9);
break;
}
}
}
// set default solver parameters
void mj_defaultSolRefImp(mjtNum* solref, mjtNum* solimp) {
if (solref) {
solref[0] = 0.02; // timeconst
solref[1] = 1; // dampratio
}
if (solimp) {
solimp[0] = 0.9; // dmin
solimp[1] = 0.95; // dmax
solimp[2] = 0.001; // width
solimp[3] = 0.5; // midpoint
solimp[4] = 2; // power
}
}
+139
View File
@@ -0,0 +1,139 @@
// Copyright 2025 DeepMind Technologies Limited
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#ifndef MUJOCO_SRC_ENGINE_ENGINE_CORE_UTIL_H_
#define MUJOCO_SRC_ENGINE_ENGINE_CORE_UTIL_H_
#include <mujoco/mjdata.h>
#include <mujoco/mjexport.h>
#include <mujoco/mjmodel.h>
#include "engine/engine_memory.h"
#ifdef __cplusplus
extern "C" {
#endif
// determine type of friction cone
MJAPI int mj_isPyramidal(const mjModel* m);
// determine type of constraint Jacobian
MJAPI int mj_isSparse(const mjModel* m);
//-------------------------- sparse chains ---------------------------------------------------------
// merge dof chains for two bodies
int mj_mergeChain(const mjModel* m, int* chain, int b1, int b2);
// merge dof chains for two simple bodies
int mj_mergeChainSimple(const mjModel* m, int* chain, int b1, int b2);
// get body chain
int mj_bodyChain(const mjModel* m, int body, int* chain);
//-------------------------- Jacobians -------------------------------------------------------------
// compute 3/6-by-nv Jacobian of global point attached to given body
MJAPI void mj_jac(const mjModel* m, const mjData* d,
mjtNum* jacp, mjtNum* jacr, const mjtNum point[3], int body);
// compute body frame Jacobian
MJAPI void mj_jacBody(const mjModel* m, const mjData* d,
mjtNum* jacp, mjtNum* jacr, int body);
// compute body center-of-mass Jacobian
MJAPI void mj_jacBodyCom(const mjModel* m, const mjData* d,
mjtNum* jacp, mjtNum* jacr, int body);
// compute subtree center-of-mass Jacobian
MJAPI void mj_jacSubtreeCom(const mjModel* m, mjData* d, mjtNum* jacp, int body);
// compute geom Jacobian
MJAPI void mj_jacGeom(const mjModel* m, const mjData* d,
mjtNum* jacp, mjtNum* jacr, int geom);
// compute site Jacobian
MJAPI void mj_jacSite(const mjModel* m, const mjData* d,
mjtNum* jacp, mjtNum* jacr, int site);
// compute translation Jacobian of point, and rotation Jacobian of axis
MJAPI void mj_jacPointAxis(const mjModel* m, mjData* d,
mjtNum* jacPoint, mjtNum* jacAxis,
const mjtNum point[3], const mjtNum axis[3], int body);
// compute 3/6-by-nv sparse Jacobian of global point attached to given body
void mj_jacSparse(const mjModel* m, const mjData* d,
mjtNum* jacp, mjtNum* jacr, const mjtNum* point, int body,
int NV, const int* chain);
// sparse Jacobian difference for simple body contacts
void mj_jacSparseSimple(const mjModel* m, const mjData* d,
mjtNum* jacdifp, mjtNum* jacdifr, const mjtNum* point,
int body, int flg_second, int NV, int start);
// dense or sparse Jacobian difference for two body points: pos2 - pos1, global
MJAPI int mj_jacDifPair(const mjModel* m, const mjData* d, int* chain,
int b1, int b2, const mjtNum pos1[3], const mjtNum pos2[3],
mjtNum* jac1p, mjtNum* jac2p, mjtNum* jacdifp,
mjtNum* jac1r, mjtNum* jac2r, mjtNum* jacdifr);
// dense or sparse weighted sum of multiple body Jacobians at same point
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 3/6-by-nv Jacobian time derivative of global point attached to given body
MJAPI void mj_jacDot(const mjModel* m, const mjData* d,
mjtNum* jacp, mjtNum* jacr, const mjtNum point[3], int body);
// compute subtree angular momentum matrix
MJAPI void mj_angmomMat(const mjModel* m, mjData* d, mjtNum* mat, int body);
//-------------------------- coordinate transformation ---------------------------------------------
// compute object 6D velocity in object-centered frame, world/local orientation
MJAPI void mj_objectVelocity(const mjModel* m, const mjData* d,
int objtype, int objid, mjtNum res[6], int flg_local);
// compute object 6D acceleration in object-centered frame, world/local orientation
MJAPI void mj_objectAcceleration(const mjModel* m, const mjData* d,
int objtype, int objid, mjtNum res[6], int flg_local);
//-------------------------- miscellaneous ---------------------------------------------------------
// map from body local to global Cartesian coordinates
MJAPI void mj_local2Global(mjData* d, mjtNum xpos[3], mjtNum xmat[9],
const mjtNum pos[3], const mjtNum quat[4],
int body, mjtByte sameframe);
// extract 6D force:torque for one contact, in contact frame
MJAPI void mj_contactForce(const mjModel* m, const mjData* d, int id, mjtNum result[6]);
// high-level warning function: count warnings in mjData, print only the first time
MJAPI void mj_warning(mjData* d, int warning, int info);
// extract 6D force:torque for one contact, in contact frame
MJAPI void mj_contactForce(const mjModel* m, const mjData* d, int id, mjtNum result[6]);
// set default solver parameters
MJAPI void mj_defaultSolRefImp(mjtNum* solref, mjtNum* solimp);
#ifdef __cplusplus
}
#endif
#endif // MUJOCO_SRC_ENGINE_ENGINE_CORE_UTIL_H_
+2 -1
View File
@@ -18,8 +18,9 @@
#include <mujoco/mjmodel.h>
#include <mujoco/mjsan.h> // IWYU pragma: keep
#include "engine/engine_core_constraint.h"
#include "engine/engine_core_util.h"
#include "engine/engine_crossplatform.h"
#include "engine/engine_io.h"
#include "engine/engine_memory.h"
#include "engine/engine_passive.h"
#include "engine/engine_support.h"
#include "engine/engine_util_blas.h"
+1
View File
@@ -24,6 +24,7 @@
#include "engine/engine_forward.h"
#include "engine/engine_io.h"
#include "engine/engine_inverse.h"
#include "engine/engine_memory.h"
#include "engine/engine_macro.h"
#include "engine/engine_support.h"
#include "engine/engine_util_blas.h"
+1
View File
@@ -31,6 +31,7 @@
#include "engine/engine_island.h"
#include "engine/engine_io.h"
#include "engine/engine_macro.h"
#include "engine/engine_memory.h"
#include "engine/engine_passive.h"
#include "engine/engine_plugin.h"
#include "engine/engine_sensor.h"
+1
View File
@@ -25,6 +25,7 @@
#include "engine/engine_core_smooth.h"
#include "engine/engine_derivative.h"
#include "engine/engine_io.h"
#include "engine/engine_memory.h"
#include "engine/engine_macro.h"
#include "engine/engine_forward.h"
#include "engine/engine_sensor.h"
+2 -361
View File
@@ -27,13 +27,14 @@
#include <mujoco/mjplugin.h>
#include <mujoco/mjsan.h> // IWYU pragma: keep
#include <mujoco/mjxmacro.h>
#include "engine/engine_core_util.h"
#include "engine/engine_crossplatform.h"
#include "engine/engine_macro.h"
#include "engine/engine_memory.h"
#include "engine/engine_plugin.h"
#include "engine/engine_util_blas.h"
#include "engine/engine_util_errmem.h"
#include "engine/engine_util_misc.h"
#include "thread/thread_pool.h"
#ifdef ADDRESS_SANITIZER
#include <sanitizer/asan_interface.h>
@@ -48,29 +49,8 @@
#pragma warning (disable: 4305) // disable MSVC warning: truncation from 'double' to 'float'
#endif
// add red zone padding when built with asan, to detect out-of-bound accesses
#ifdef ADDRESS_SANITIZER
#define mjREDZONE 32
#else
#define mjREDZONE 0
#endif
static const int MAX_ARRAY_SIZE = INT_MAX / 4;
// compute a % b with a fast code path if the second argument is a power of 2
static inline size_t fastmod(size_t a, size_t b) {
// (b & (b - 1)) == 0 implies that b is a power of 2
if (mjLIKELY((b & (b - 1)) == 0)) {
return a & (b - 1);
}
return a % b;
}
typedef struct {
size_t pbase; // value of d->pbase immediately before mj_markStack
size_t pstack; // value of d->pstack immediately before mj_markStack
void* pc; // program counter of the call site of mj_markStack (only set when under asan)
} mjStackFrame;
//------------------------------ mjLROpt -----------------------------------------------------------
@@ -93,22 +73,6 @@ void mj_defaultLROpt(mjLROpt* opt) {
//------------------------------- mjOption ---------------------------------------------------------
// set default solver parameters
void mj_defaultSolRefImp(mjtNum* solref, mjtNum* solimp) {
if (solref) {
solref[0] = 0.02; // timeconst
solref[1] = 1; // dampratio
}
if (solimp) {
solimp[0] = 0.9; // dmin
solimp[1] = 0.95; // dmax
solimp[2] = 0.001; // width
solimp[3] = 0.5; // midpoint
solimp[4] = 2; // power
}
}
// set model options to default values
@@ -1513,329 +1477,6 @@ mjData* mjv_copyData(mjData* dest, const mjModel* m, const mjData* src) {
return mj_copyDataVisual(dest, m, src, /*flg_all=*/0);
}
static void maybe_lock_alloc_mutex(mjData* d) {
if (d->threadpool != 0) {
mju_threadPoolLockAllocMutex((mjThreadPool*)d->threadpool);
}
}
static void maybe_unlock_alloc_mutex(mjData* d) {
if (d->threadpool != 0) {
mju_threadPoolUnlockAllocMutex((mjThreadPool*)d->threadpool);
}
}
static inline mjStackInfo get_stack_info_from_data(const mjData* d) {
mjStackInfo stack_info;
stack_info.bottom = (uintptr_t)d->arena + (uintptr_t)d->narena;
stack_info.top = stack_info.bottom - d->pstack;
stack_info.limit = (uintptr_t)d->arena + (uintptr_t)d->parena;
stack_info.stack_base = d->pbase;
return stack_info;
}
#ifdef ADDRESS_SANITIZER
// get stack usage from red-zone (under ASAN)
static size_t stack_usage_redzone(const mjStackInfo* stack_info) {
size_t usage = 0;
// actual stack usage (without red zone bytes) is stored in the red zone
if (stack_info->top != stack_info->bottom) {
char* prev_pstack_ptr = (char*)(stack_info->top);
size_t prev_misalign = (uintptr_t)prev_pstack_ptr % _Alignof(size_t);
size_t* prev_usage_ptr =
(size_t*)(prev_pstack_ptr +
(prev_misalign ? _Alignof(size_t) - prev_misalign : 0));
ASAN_UNPOISON_MEMORY_REGION(prev_usage_ptr, sizeof(size_t));
usage = *prev_usage_ptr;
ASAN_POISON_MEMORY_REGION(prev_usage_ptr, sizeof(size_t));
}
return usage;
}
#endif
// allocate memory from the mjData arena
void* mj_arenaAllocByte(mjData* d, size_t bytes, size_t alignment) {
maybe_lock_alloc_mutex(d);
size_t misalignment = fastmod(d->parena, alignment);
size_t padding = misalignment ? alignment - misalignment : 0;
// check size
size_t bytes_available = d->narena - d->pstack;
if (mjUNLIKELY(d->parena + padding + bytes > bytes_available)) {
maybe_unlock_alloc_mutex(d);
return NULL;
}
size_t stack_usage = d->pstack;
// under ASAN, get stack usage from red zone
#ifdef ADDRESS_SANITIZER
mjStackInfo stack_info;
mjStackInfo* stack_info_ptr;
if (!d->threadpool) {
stack_info = get_stack_info_from_data(d);
stack_info_ptr = &stack_info;
} else {
size_t thread_id = mju_threadPoolCurrentWorkerId((mjThreadPool*)d->threadpool);
stack_info_ptr = mju_getStackInfoForThread(d, thread_id);
}
stack_usage = stack_usage_redzone(stack_info_ptr);
#endif
// allocate, update max, return pointer to buffer
void* result = (char*)d->arena + d->parena + padding;
d->parena += padding + bytes;
d->maxuse_arena = mjMAX(d->maxuse_arena, stack_usage + d->parena);
#ifdef ADDRESS_SANITIZER
ASAN_UNPOISON_MEMORY_REGION(result, bytes);
#endif
#ifdef MEMORY_SANITIZER
__msan_allocated_memory(result, bytes);
#endif
maybe_unlock_alloc_mutex(d);
return result;
}
// internal: allocate size bytes on the provided stack shard
// declared inline so that modular arithmetic with specific alignments can be optimized out
static inline void* stackallocinternal(mjData* d, mjStackInfo* stack_info, size_t size,
size_t alignment, const char* caller, int line) {
// return NULL if empty
if (mjUNLIKELY(!size)) {
return NULL;
}
// start of the memory to be allocated to the buffer
uintptr_t start_ptr = stack_info->top - (size + mjREDZONE);
// align the pointer
start_ptr -= fastmod(start_ptr, alignment);
// new top of the stack
uintptr_t new_top_ptr = start_ptr - mjREDZONE;
// exclude red zone from stack usage statistics
size_t current_alloc_usage = stack_info->top - new_top_ptr - 2 * mjREDZONE;
size_t usage = current_alloc_usage + (stack_info->bottom - stack_info->top);
// check size
size_t stack_available_bytes = stack_info->top - stack_info->limit;
size_t stack_required_bytes = stack_info->top - new_top_ptr;
if (mjUNLIKELY(stack_required_bytes > stack_available_bytes)) {
char info[1024];
if (caller) {
snprintf(info, sizeof(info), " at %s, line %d", caller, line);
} else {
info[0] = '\0';
}
mju_error(
"mj_stackAlloc: out of memory, stack overflow%s\n"
" max = %" PRIuPTR ", available = %" PRIuPTR ", requested = %" PRIuPTR
"\n nefc = %d, ncon = %d",
info, stack_info->bottom - stack_info->limit, stack_available_bytes,
stack_required_bytes, d->nefc, d->ncon);
}
#ifdef ADDRESS_SANITIZER
usage = current_alloc_usage + stack_usage_redzone(stack_info);
// store new stack usage in the red zone
size_t misalign = new_top_ptr % _Alignof(size_t);
size_t* usage_ptr =
(size_t*)(new_top_ptr + (misalign ? _Alignof(size_t) - misalign : 0));
ASAN_UNPOISON_MEMORY_REGION(usage_ptr, sizeof(size_t));
*usage_ptr = usage;
ASAN_POISON_MEMORY_REGION(usage_ptr, sizeof(size_t));
// unpoison the actual usable allocation
ASAN_UNPOISON_MEMORY_REGION((void*)start_ptr, size);
#endif
// update max usage statistics
stack_info->top = new_top_ptr;
if (!d->threadpool) {
d->maxuse_stack = mjMAX(d->maxuse_stack, usage);
d->maxuse_arena = mjMAX(d->maxuse_arena, usage + d->parena);
} else {
size_t thread_id = mju_threadPoolCurrentWorkerId((mjThreadPool*)d->threadpool);
d->maxuse_threadstack[thread_id] = mjMAX(d->maxuse_threadstack[thread_id], usage);
}
return (void*)start_ptr;
}
// internal: allocate size bytes in mjData
// declared inline so that modular arithmetic with specific alignments can be optimized out
static inline void* stackalloc(mjData* d, size_t size, size_t alignment,
const char* caller, int line) {
// single threaded allocation
if (!d->threadpool) {
mjStackInfo stack_info = get_stack_info_from_data(d);
void* result = stackallocinternal(d, &stack_info, size, alignment, caller, line);
d->pstack = stack_info.bottom - stack_info.top;
return result;
}
// multi threaded allocation
size_t thread_id = mju_threadPoolCurrentWorkerId((mjThreadPool*)d->threadpool);
mjStackInfo* stack_info = mju_getStackInfoForThread(d, thread_id);
return stackallocinternal(d, stack_info, size, alignment, caller, line);
}
// mjStackInfo mark stack frame, inline so ASAN errors point to correct code unit
#ifdef ADDRESS_SANITIZER
__attribute__((always_inline))
#endif
static inline void markstackinternal(mjData* d, mjStackInfo* stack_info) {
size_t top_old = stack_info->top;
mjStackFrame* s =
(mjStackFrame*) stackallocinternal(d, stack_info, sizeof(mjStackFrame), _Alignof(mjStackFrame), NULL, 0);
s->pbase = stack_info->stack_base;
s->pstack = top_old;
#ifdef ADDRESS_SANITIZER
// store the program counter to the caller so that we can compare against mj_freeStack later
s->pc = __sanitizer_return_address();
#endif
stack_info->stack_base = (uintptr_t) s;
}
// mjData mark stack frame
#ifndef ADDRESS_SANITIZER
void mj_markStack(mjData* d)
#else
void mj__markStack(mjData* d)
#endif
{
if (!d->threadpool) {
mjStackInfo stack_info = get_stack_info_from_data(d);
markstackinternal(d, &stack_info);
d->pstack = stack_info.bottom - stack_info.top;
d->pbase = stack_info.stack_base;
return;
}
size_t thread_id = mju_threadPoolCurrentWorkerId((mjThreadPool*)d->threadpool);
mjStackInfo* stack_info = mju_getStackInfoForThread(d, thread_id);
markstackinternal(d, stack_info);
}
#ifdef ADDRESS_SANITIZER
__attribute__((always_inline))
#endif
static inline void freestackinternal(mjStackInfo* stack_info) {
if (mjUNLIKELY(!stack_info->stack_base)) {
return;
}
mjStackFrame* s = (mjStackFrame*) stack_info->stack_base;
#ifdef ADDRESS_SANITIZER
// raise an error if caller function name doesn't match the most recent caller of mj_markStack
if (!mj__comparePcFuncName(s->pc, __sanitizer_return_address())) {
mjERROR("mj_markStack %s has no corresponding mj_freeStack (detected %s)",
mj__getPcDebugInfo(s->pc),
mj__getPcDebugInfo(__sanitizer_return_address()));
}
#endif
// restore pbase and pstack
stack_info->stack_base = s->pbase;
stack_info->top = s->pstack;
// if running under asan, poison the newly freed memory region
#ifdef ADDRESS_SANITIZER
ASAN_POISON_MEMORY_REGION((char*)stack_info->limit, stack_info->top - stack_info->limit);
#endif
}
// mjData free stack frame
#ifndef ADDRESS_SANITIZER
void mj_freeStack(mjData* d)
#else
void mj__freeStack(mjData* d)
#endif
{
if (!d->threadpool) {
mjStackInfo stack_info = get_stack_info_from_data(d);
freestackinternal(&stack_info);
d->pstack = stack_info.bottom - stack_info.top;
d->pbase = stack_info.stack_base;
return;
}
size_t thread_id = mju_threadPoolCurrentWorkerId((mjThreadPool*)d->threadpool);
mjStackInfo* stack_info = mju_getStackInfoForThread(d, thread_id);
freestackinternal(stack_info);
}
// returns the number of bytes available on the stack
size_t mj_stackBytesAvailable(mjData* d) {
if (!d->threadpool) {
mjStackInfo stack_info = get_stack_info_from_data(d);
return stack_info.top - stack_info.limit;
} else {
size_t thread_id = mju_threadPoolCurrentWorkerId((mjThreadPool*)d->threadpool);
mjStackInfo* stack_info = mju_getStackInfoForThread(d, thread_id);
return stack_info->top - stack_info->limit;
}
}
// allocate bytes on the stack
void* mj_stackAllocByte(mjData* d, size_t bytes, size_t alignment) {
return stackalloc(d, bytes, alignment, NULL, 0);
}
// allocate bytes on the stack, with caller information
void* mj_stackAllocInfo(mjData* d, size_t bytes, size_t alignment,
const char* caller, int line) {
return stackalloc(d, bytes, alignment, caller, line);
}
// allocate mjtNums on the stack
mjtNum* mj_stackAllocNum(mjData* d, size_t size) {
if (mjUNLIKELY(size >= SIZE_MAX / sizeof(mjtNum))) {
mjERROR("requested size is too large (more than 2^64 bytes).");
}
return (mjtNum*) stackalloc(d, size * sizeof(mjtNum), _Alignof(mjtNum), NULL, 0);
}
// allocate ints on the stack
int* mj_stackAllocInt(mjData* d, size_t size) {
if (mjUNLIKELY(size >= SIZE_MAX / sizeof(int))) {
mjERROR("requested size is too large (more than 2^64 bytes).");
}
return (int*) stackalloc(d, size * sizeof(int), _Alignof(int), NULL, 0);
}
// clear data, set defaults
-51
View File
@@ -35,9 +35,6 @@ extern "C" {
// Set default options for length range computation.
MJAPI void mj_defaultLROpt(mjLROpt* opt);
// set default solver parameters
MJAPI void mj_defaultSolRefImp(mjtNum* solref, mjtNum* solimp);
// set options to default values
MJAPI void mj_defaultOption(mjOption* opt);
@@ -130,60 +127,12 @@ MJAPI void mj_resetDataDebug(const mjModel* m, mjData* d, unsigned char debug_va
// Reset data. If 0 <= key < nkey, set fields from specified keyframe.
MJAPI void mj_resetDataKeyframe(const mjModel* m, mjData* d, int key);
// mjData arena allocate
MJAPI void* mj_arenaAllocByte(mjData* d, size_t bytes, size_t alignment);
// init plugins
MJAPI void mj_initPlugin(const mjModel* m, mjData* d);
#ifndef ADDRESS_SANITIZER
// mjData mark stack frame
MJAPI void mj_markStack(mjData* d);
// mjData free stack frame
MJAPI void mj_freeStack(mjData* d);
#else
void mj__markStack(mjData* d) __attribute__((noinline));
void mj__freeStack(mjData* d) __attribute__((noinline));
#endif // ADDRESS_SANITIZER
// returns the number of bytes available on the stack
MJAPI size_t mj_stackBytesAvailable(mjData* d);
// allocate bytes on the stack
MJAPI void* mj_stackAllocByte(mjData* d, size_t bytes, size_t alignment);
// allocate bytes on the stack, with added caller information
MJAPI void* mj_stackAllocInfo(mjData* d, size_t bytes, size_t alignment,
const char* caller, int line);
// macro to allocate a stack array of given type, adds caller information
#define mjSTACKALLOC(d, num, type) \
(type*) mj_stackAllocInfo(d, (num) * sizeof(type), _Alignof(type), __func__, __LINE__)
// mjData stack allocate for array of mjtNums
MJAPI mjtNum* mj_stackAllocNum(mjData* d, size_t size);
// mjData stack allocate for array of ints
MJAPI int* mj_stackAllocInt(mjData* d, size_t size);
// deallocate data
MJAPI void mj_deleteData(mjData* d);
// clear arena pointers in mjData
static inline void mj_clearEfc(mjData* d) {
#define X(type, name, nr, nc) d->name = NULL;
MJDATA_ARENA_POINTERS
#undef X
d->nefc = 0;
d->nisland = 0;
d->contact = (mjContact*) d->arena;
}
#ifdef __cplusplus
}
#endif
+2
View File
@@ -23,7 +23,9 @@
#include <mujoco/mjsan.h> // IWYU pragma: keep
#include <mujoco/mjxmacro.h>
#include "engine/engine_core_constraint.h"
#include "engine/engine_core_util.h"
#include "engine/engine_io.h"
#include "engine/engine_memory.h"
#include "engine/engine_support.h"
#include "engine/engine_util_errmem.h"
#include "engine/engine_util_misc.h"
+394
View File
@@ -0,0 +1,394 @@
// Copyright 2025 DeepMind Technologies Limited
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#include "engine/engine_memory.h"
#include <inttypes.h> // NOLINT required for PRIu64, PRIuPTR
#include <limits.h>
#include <stddef.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <mujoco/mjmacro.h>
#include <mujoco/mjmodel.h>
#include <mujoco/mjplugin.h>
#include <mujoco/mjsan.h> // IWYU pragma: keep
#include <mujoco/mjxmacro.h>
#include "engine/engine_crossplatform.h"
#include "engine/engine_macro.h"
#include "engine/engine_plugin.h"
#include "engine/engine_util_blas.h"
#include "engine/engine_util_errmem.h"
#include "engine/engine_util_misc.h"
#include "thread/thread_pool.h"
#ifdef ADDRESS_SANITIZER
#include <sanitizer/asan_interface.h>
#include <sanitizer/common_interface_defs.h>
#endif
#ifdef MEMORY_SANITIZER
#include <sanitizer/msan_interface.h>
#endif
#ifdef _MSC_VER
#pragma warning (disable: 4305) // disable MSVC warning: truncation from 'double' to 'float'
#endif
// add red zone padding when built with asan, to detect out-of-bound accesses
#ifdef ADDRESS_SANITIZER
#define mjREDZONE 32
#else
#define mjREDZONE 0
#endif
// compute a % b with a fast code path if the second argument is a power of 2
static inline size_t fastmod(size_t a, size_t b) {
// (b & (b - 1)) == 0 implies that b is a power of 2
if (mjLIKELY((b & (b - 1)) == 0)) {
return a & (b - 1);
}
return a % b;
}
typedef struct {
size_t pbase; // value of d->pbase immediately before mj_markStack
size_t pstack; // value of d->pstack immediately before mj_markStack
void* pc; // program counter of the call site of mj_markStack (only set when under asan)
} mjStackFrame;
static void maybe_lock_alloc_mutex(mjData* d) {
if (d->threadpool != 0) {
mju_threadPoolLockAllocMutex((mjThreadPool*)d->threadpool);
}
}
static void maybe_unlock_alloc_mutex(mjData* d) {
if (d->threadpool != 0) {
mju_threadPoolUnlockAllocMutex((mjThreadPool*)d->threadpool);
}
}
static inline mjStackInfo get_stack_info_from_data(const mjData* d) {
mjStackInfo stack_info;
stack_info.bottom = (uintptr_t)d->arena + (uintptr_t)d->narena;
stack_info.top = stack_info.bottom - d->pstack;
stack_info.limit = (uintptr_t)d->arena + (uintptr_t)d->parena;
stack_info.stack_base = d->pbase;
return stack_info;
}
#ifdef ADDRESS_SANITIZER
// get stack usage from red-zone (under ASAN)
static size_t stack_usage_redzone(const mjStackInfo* stack_info) {
size_t usage = 0;
// actual stack usage (without red zone bytes) is stored in the red zone
if (stack_info->top != stack_info->bottom) {
char* prev_pstack_ptr = (char*)(stack_info->top);
size_t prev_misalign = (uintptr_t)prev_pstack_ptr % _Alignof(size_t);
size_t* prev_usage_ptr =
(size_t*)(prev_pstack_ptr +
(prev_misalign ? _Alignof(size_t) - prev_misalign : 0));
ASAN_UNPOISON_MEMORY_REGION(prev_usage_ptr, sizeof(size_t));
usage = *prev_usage_ptr;
ASAN_POISON_MEMORY_REGION(prev_usage_ptr, sizeof(size_t));
}
return usage;
}
#endif
// allocate memory from the mjData arena
void* mj_arenaAllocByte(mjData* d, size_t bytes, size_t alignment) {
maybe_lock_alloc_mutex(d);
size_t misalignment = fastmod(d->parena, alignment);
size_t padding = misalignment ? alignment - misalignment : 0;
// check size
size_t bytes_available = d->narena - d->pstack;
if (mjUNLIKELY(d->parena + padding + bytes > bytes_available)) {
maybe_unlock_alloc_mutex(d);
return NULL;
}
size_t stack_usage = d->pstack;
// under ASAN, get stack usage from red zone
#ifdef ADDRESS_SANITIZER
mjStackInfo stack_info;
mjStackInfo* stack_info_ptr;
if (!d->threadpool) {
stack_info = get_stack_info_from_data(d);
stack_info_ptr = &stack_info;
} else {
size_t thread_id = mju_threadPoolCurrentWorkerId((mjThreadPool*)d->threadpool);
stack_info_ptr = mju_getStackInfoForThread(d, thread_id);
}
stack_usage = stack_usage_redzone(stack_info_ptr);
#endif
// allocate, update max, return pointer to buffer
void* result = (char*)d->arena + d->parena + padding;
d->parena += padding + bytes;
d->maxuse_arena = mjMAX(d->maxuse_arena, stack_usage + d->parena);
#ifdef ADDRESS_SANITIZER
ASAN_UNPOISON_MEMORY_REGION(result, bytes);
#endif
#ifdef MEMORY_SANITIZER
__msan_allocated_memory(result, bytes);
#endif
maybe_unlock_alloc_mutex(d);
return result;
}
// internal: allocate size bytes on the provided stack shard
// declared inline so that modular arithmetic with specific alignments can be optimized out
static inline void* stackallocinternal(mjData* d, mjStackInfo* stack_info, size_t size,
size_t alignment, const char* caller, int line) {
// return NULL if empty
if (mjUNLIKELY(!size)) {
return NULL;
}
// start of the memory to be allocated to the buffer
uintptr_t start_ptr = stack_info->top - (size + mjREDZONE);
// align the pointer
start_ptr -= fastmod(start_ptr, alignment);
// new top of the stack
uintptr_t new_top_ptr = start_ptr - mjREDZONE;
// exclude red zone from stack usage statistics
size_t current_alloc_usage = stack_info->top - new_top_ptr - 2 * mjREDZONE;
size_t usage = current_alloc_usage + (stack_info->bottom - stack_info->top);
// check size
size_t stack_available_bytes = stack_info->top - stack_info->limit;
size_t stack_required_bytes = stack_info->top - new_top_ptr;
if (mjUNLIKELY(stack_required_bytes > stack_available_bytes)) {
char info[1024];
if (caller) {
snprintf(info, sizeof(info), " at %s, line %d", caller, line);
} else {
info[0] = '\0';
}
mju_error(
"mj_stackAlloc: out of memory, stack overflow%s\n"
" max = %" PRIuPTR ", available = %" PRIuPTR ", requested = %" PRIuPTR
"\n nefc = %d, ncon = %d",
info, stack_info->bottom - stack_info->limit, stack_available_bytes,
stack_required_bytes, d->nefc, d->ncon);
}
#ifdef ADDRESS_SANITIZER
usage = current_alloc_usage + stack_usage_redzone(stack_info);
// store new stack usage in the red zone
size_t misalign = new_top_ptr % _Alignof(size_t);
size_t* usage_ptr =
(size_t*)(new_top_ptr + (misalign ? _Alignof(size_t) - misalign : 0));
ASAN_UNPOISON_MEMORY_REGION(usage_ptr, sizeof(size_t));
*usage_ptr = usage;
ASAN_POISON_MEMORY_REGION(usage_ptr, sizeof(size_t));
// unpoison the actual usable allocation
ASAN_UNPOISON_MEMORY_REGION((void*)start_ptr, size);
#endif
// update max usage statistics
stack_info->top = new_top_ptr;
if (!d->threadpool) {
d->maxuse_stack = mjMAX(d->maxuse_stack, usage);
d->maxuse_arena = mjMAX(d->maxuse_arena, usage + d->parena);
} else {
size_t thread_id = mju_threadPoolCurrentWorkerId((mjThreadPool*)d->threadpool);
d->maxuse_threadstack[thread_id] = mjMAX(d->maxuse_threadstack[thread_id], usage);
}
return (void*)start_ptr;
}
// internal: allocate size bytes in mjData
// declared inline so that modular arithmetic with specific alignments can be optimized out
static inline void* stackalloc(mjData* d, size_t size, size_t alignment,
const char* caller, int line) {
// single threaded allocation
if (!d->threadpool) {
mjStackInfo stack_info = get_stack_info_from_data(d);
void* result = stackallocinternal(d, &stack_info, size, alignment, caller, line);
d->pstack = stack_info.bottom - stack_info.top;
return result;
}
// multi threaded allocation
size_t thread_id = mju_threadPoolCurrentWorkerId((mjThreadPool*)d->threadpool);
mjStackInfo* stack_info = mju_getStackInfoForThread(d, thread_id);
return stackallocinternal(d, stack_info, size, alignment, caller, line);
}
// mjStackInfo mark stack frame, inline so ASAN errors point to correct code unit
#ifdef ADDRESS_SANITIZER
__attribute__((always_inline))
#endif
static inline void markstackinternal(mjData* d, mjStackInfo* stack_info) {
size_t top_old = stack_info->top;
mjStackFrame* s =
(mjStackFrame*) stackallocinternal(d, stack_info, sizeof(mjStackFrame), _Alignof(mjStackFrame), NULL, 0);
s->pbase = stack_info->stack_base;
s->pstack = top_old;
#ifdef ADDRESS_SANITIZER
// store the program counter to the caller so that we can compare against mj_freeStack later
s->pc = __sanitizer_return_address();
#endif
stack_info->stack_base = (uintptr_t) s;
}
// mjData mark stack frame
#ifndef ADDRESS_SANITIZER
void mj_markStack(mjData* d)
#else
void mj__markStack(mjData* d)
#endif
{
if (!d->threadpool) {
mjStackInfo stack_info = get_stack_info_from_data(d);
markstackinternal(d, &stack_info);
d->pstack = stack_info.bottom - stack_info.top;
d->pbase = stack_info.stack_base;
return;
}
size_t thread_id = mju_threadPoolCurrentWorkerId((mjThreadPool*)d->threadpool);
mjStackInfo* stack_info = mju_getStackInfoForThread(d, thread_id);
markstackinternal(d, stack_info);
}
#ifdef ADDRESS_SANITIZER
__attribute__((always_inline))
#endif
static inline void freestackinternal(mjStackInfo* stack_info) {
if (mjUNLIKELY(!stack_info->stack_base)) {
return;
}
mjStackFrame* s = (mjStackFrame*) stack_info->stack_base;
#ifdef ADDRESS_SANITIZER
// raise an error if caller function name doesn't match the most recent caller of mj_markStack
if (!mj__comparePcFuncName(s->pc, __sanitizer_return_address())) {
mjERROR("mj_markStack %s has no corresponding mj_freeStack (detected %s)",
mj__getPcDebugInfo(s->pc),
mj__getPcDebugInfo(__sanitizer_return_address()));
}
#endif
// restore pbase and pstack
stack_info->stack_base = s->pbase;
stack_info->top = s->pstack;
// if running under asan, poison the newly freed memory region
#ifdef ADDRESS_SANITIZER
ASAN_POISON_MEMORY_REGION((char*)stack_info->limit, stack_info->top - stack_info->limit);
#endif
}
// mjData free stack frame
#ifndef ADDRESS_SANITIZER
void mj_freeStack(mjData* d)
#else
void mj__freeStack(mjData* d)
#endif
{
if (!d->threadpool) {
mjStackInfo stack_info = get_stack_info_from_data(d);
freestackinternal(&stack_info);
d->pstack = stack_info.bottom - stack_info.top;
d->pbase = stack_info.stack_base;
return;
}
size_t thread_id = mju_threadPoolCurrentWorkerId((mjThreadPool*)d->threadpool);
mjStackInfo* stack_info = mju_getStackInfoForThread(d, thread_id);
freestackinternal(stack_info);
}
// returns the number of bytes available on the stack
size_t mj_stackBytesAvailable(mjData* d) {
if (!d->threadpool) {
mjStackInfo stack_info = get_stack_info_from_data(d);
return stack_info.top - stack_info.limit;
} else {
size_t thread_id = mju_threadPoolCurrentWorkerId((mjThreadPool*)d->threadpool);
mjStackInfo* stack_info = mju_getStackInfoForThread(d, thread_id);
return stack_info->top - stack_info->limit;
}
}
// allocate bytes on the stack
void* mj_stackAllocByte(mjData* d, size_t bytes, size_t alignment) {
return stackalloc(d, bytes, alignment, NULL, 0);
}
// allocate bytes on the stack, with caller information
void* mj_stackAllocInfo(mjData* d, size_t bytes, size_t alignment,
const char* caller, int line) {
return stackalloc(d, bytes, alignment, caller, line);
}
// allocate mjtNums on the stack
mjtNum* mj_stackAllocNum(mjData* d, size_t size) {
if (mjUNLIKELY(size >= SIZE_MAX / sizeof(mjtNum))) {
mjERROR("requested size is too large (more than 2^64 bytes).");
}
return (mjtNum*) stackalloc(d, size * sizeof(mjtNum), _Alignof(mjtNum), NULL, 0);
}
// allocate ints on the stack
int* mj_stackAllocInt(mjData* d, size_t size) {
if (mjUNLIKELY(size >= SIZE_MAX / sizeof(int))) {
mjERROR("requested size is too large (more than 2^64 bytes).");
}
return (int*) stackalloc(d, size * sizeof(int), _Alignof(int), NULL, 0);
}
+85
View File
@@ -0,0 +1,85 @@
// Copyright 2025 DeepMind Technologies Limited
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#ifndef MUJOCO_SRC_ENGINE_ENGINE_MEMORY_H_
#define MUJOCO_SRC_ENGINE_ENGINE_MEMORY_H_
#include <mujoco/mjdata.h>
#include <mujoco/mjexport.h>
#include <mujoco/mjmodel.h>
#include <mujoco/mjxmacro.h>
#ifdef __cplusplus
#include <cstddef>
extern "C" {
#else
#include <stddef.h>
#endif
// internal hash map size factor (2 corresponds to a load factor of 0.5)
#define mjLOAD_MULTIPLE 2
// mjData arena allocate
MJAPI void* mj_arenaAllocByte(mjData* d, size_t bytes, size_t alignment);
#ifndef ADDRESS_SANITIZER
// mjData mark stack frame
MJAPI void mj_markStack(mjData* d);
// mjData free stack frame
MJAPI void mj_freeStack(mjData* d);
#else
void mj__markStack(mjData* d) __attribute__((noinline));
void mj__freeStack(mjData* d) __attribute__((noinline));
#endif // ADDRESS_SANITIZER
// returns the number of bytes available on the stack
MJAPI size_t mj_stackBytesAvailable(mjData* d);
// allocate bytes on the stack
MJAPI void* mj_stackAllocByte(mjData* d, size_t bytes, size_t alignment);
// allocate bytes on the stack, with added caller information
MJAPI void* mj_stackAllocInfo(mjData* d, size_t bytes, size_t alignment,
const char* caller, int line);
// macro to allocate a stack array of given type, adds caller information
#define mjSTACKALLOC(d, num, type) \
(type*) mj_stackAllocInfo(d, (num) * sizeof(type), _Alignof(type), __func__, __LINE__)
// mjData stack allocate for array of mjtNums
MJAPI mjtNum* mj_stackAllocNum(mjData* d, size_t size);
// mjData stack allocate for array of ints
MJAPI int* mj_stackAllocInt(mjData* d, size_t size);
// clear arena pointers in mjData
static inline void mj_clearEfc(mjData* d) {
#define X(type, name, nr, nc) d->name = NULL;
MJDATA_ARENA_POINTERS
#undef X
d->nefc = 0;
d->nisland = 0;
d->contact = (mjContact*) d->arena;
}
#ifdef __cplusplus
}
#endif
#endif // MUJOCO_SRC_ENGINE_ENGINE_MEMORY_H_
+2
View File
@@ -22,8 +22,10 @@
#include <mujoco/mjmodel.h>
#include "engine/engine_callback.h"
#include "engine/engine_core_constraint.h"
#include "engine/engine_core_util.h"
#include "engine/engine_crossplatform.h"
#include "engine/engine_io.h"
#include "engine/engine_memory.h"
#include "engine/engine_plugin.h"
#include "engine/engine_support.h"
#include "engine/engine_util_blas.h"
+1
View File
@@ -25,6 +25,7 @@
#include <mujoco/mjsan.h> // IWYU pragma: keep
#include <mujoco/mjxmacro.h>
#include "engine/engine_core_constraint.h"
#include "engine/engine_core_util.h"
#include "engine/engine_io.h"
#include "engine/engine_name.h"
#include "engine/engine_macro.h"
+1
View File
@@ -24,6 +24,7 @@
#include <mujoco/mjvisualize.h>
#include "engine/engine_collision_sdf.h"
#include "engine/engine_io.h"
#include "engine/engine_memory.h"
#include "engine/engine_plugin.h"
#include "engine/engine_util_blas.h"
#include "engine/engine_util_errmem.h"
+2 -1
View File
@@ -23,8 +23,9 @@
#include "engine/engine_callback.h"
#include "engine/engine_collision_sdf.h"
#include "engine/engine_core_smooth.h"
#include "engine/engine_core_util.h"
#include "engine/engine_crossplatform.h"
#include "engine/engine_io.h"
#include "engine/engine_memory.h"
#include "engine/engine_plugin.h"
#include "engine/engine_ray.h"
#include "engine/engine_sort.h"
+2 -1
View File
@@ -23,9 +23,10 @@
#include <mujoco/mjsan.h> // IWYU pragma: keep
#include "engine/engine_core_constraint.h"
#include "engine/engine_core_smooth.h"
#include "engine/engine_core_util.h"
#include "engine/engine_forward.h"
#include "engine/engine_io.h"
#include "engine/engine_support.h"
#include "engine/engine_memory.h"
#include "engine/engine_util_blas.h"
#include "engine/engine_util_errmem.h"
#include "engine/engine_util_misc.h"
+2
View File
@@ -23,7 +23,9 @@
#include <mujoco/mjsan.h> // IWYU pragma: keep
#include "engine/engine_core_constraint.h"
#include "engine/engine_core_smooth.h"
#include "engine/engine_core_util.h"
#include "engine/engine_io.h"
#include "engine/engine_memory.h"
#include "engine/engine_support.h"
#include "engine/engine_util_blas.h"
#include "engine/engine_util_errmem.h"
+2 -894
View File
@@ -23,9 +23,9 @@
#include "engine/engine_collision_driver.h"
#include "engine/engine_collision_gjk.h"
#include "engine/engine_collision_primitive.h"
#include "engine/engine_core_constraint.h"
#include "engine/engine_core_util.h"
#include "engine/engine_crossplatform.h"
#include "engine/engine_io.h"
#include "engine/engine_memory.h"
#include "engine/engine_util_blas.h"
#include "engine/engine_util_errmem.h"
#include "engine/engine_util_misc.h"
@@ -270,700 +270,6 @@ void mj_setKeyframe(mjModel* m, const mjData* d, int k) {
//-------------------------- sparse chains ---------------------------------------------------------
// merge dof chains for two bodies
int mj_mergeChain(const mjModel* m, int* chain, int b1, int b2) {
int da1, da2, NV = 0;
// skip fixed bodies
while (b1 && !m->body_dofnum[b1]) {
b1 = m->body_parentid[b1];
}
while (b2 && !m->body_dofnum[b2]) {
b2 = m->body_parentid[b2];
}
// neither body is movable: empty chain
if (b1 == 0 && b2 == 0) {
return 0;
}
// initialize last dof address for each body
da1 = m->body_dofadr[b1] + m->body_dofnum[b1] - 1;
da2 = m->body_dofadr[b2] + m->body_dofnum[b2] - 1;
// merge chains
while (da1 >= 0 || da2 >= 0) {
chain[NV] = mjMAX(da1, da2);
if (da1 == chain[NV]) {
da1 = m->dof_parentid[da1];
}
if (da2 == chain[NV]) {
da2 = m->dof_parentid[da2];
}
NV++;
}
// reverse order of chain: make it increasing
for (int i=0; i < NV/2; i++) {
int tmp = chain[i];
chain[i] = chain[NV-i-1];
chain[NV-i-1] = tmp;
}
return NV;
}
// merge dof chains for two simple bodies
int mj_mergeChainSimple(const mjModel* m, int* chain, int b1, int b2) {
// swap bodies if wrong order
if (b1 > b2) {
int tmp = b1;
b1 = b2;
b2 = tmp;
}
// init
int n1 = m->body_dofnum[b1], n2 = m->body_dofnum[b2];
// both fixed: nothing to do
if (n1 == 0 && n2 == 0) {
return 0;
}
// copy b1 dofs
for (int i=0; i < n1; i++) {
chain[i] = m->body_dofadr[b1] + i;
}
// copy b2 dofs
for (int i=0; i < n2; i++) {
chain[n1+i] = m->body_dofadr[b2] + i;
}
return (n1+n2);
}
// get body chain
int mj_bodyChain(const mjModel* m, int body, int* chain) {
// simple body
if (m->body_simple[body]) {
int dofnum = m->body_dofnum[body];
for (int i=0; i < dofnum; i++) {
chain[i] = m->body_dofadr[body] + i;
}
return dofnum;
}
// general case
else {
// skip fixed bodies
while (body && !m->body_dofnum[body]) {
body = m->body_parentid[body];
}
// not movable: empty chain
if (body == 0) {
return 0;
}
// initialize last dof
int da = m->body_dofadr[body] + m->body_dofnum[body] - 1;
int NV = 0;
// construct chain from child to parent
while (da >= 0) {
chain[NV++] = da;
da = m->dof_parentid[da];
}
// reverse order of chain: make it increasing
for (int i=0; i < NV/2; i++) {
int tmp = chain[i];
chain[i] = chain[NV-i-1];
chain[NV-i-1] = tmp;
}
return NV;
}
}
//-------------------------- Jacobians -------------------------------------------------------------
// compute 3/6-by-nv Jacobian of global point attached to given body
void mj_jac(const mjModel* m, const mjData* d,
mjtNum* jacp, mjtNum* jacr, const mjtNum point[3], int body) {
int nv = m->nv;
mjtNum offset[3];
// clear jacobians, compute offset if required
if (jacp) {
mju_zero(jacp, 3*nv);
mju_sub3(offset, point, d->subtree_com+3*m->body_rootid[body]);
}
if (jacr) {
mju_zero(jacr, 3*nv);
}
// skip fixed bodies
while (body && !m->body_dofnum[body]) {
body = m->body_parentid[body];
}
// no movable body found: nothing to do
if (!body) {
return;
}
// get last dof that affects this (as well as the original) body
int i = m->body_dofadr[body] + m->body_dofnum[body] - 1;
// backward pass over dof ancestor chain
while (i >= 0) {
mjtNum* cdof = d->cdof+6*i;
// construct rotation jacobian
if (jacr) {
jacr[i+0*nv] = cdof[0];
jacr[i+1*nv] = cdof[1];
jacr[i+2*nv] = cdof[2];
}
// construct translation jacobian (correct for rotation)
if (jacp) {
mjtNum tmp[3];
mju_cross(tmp, cdof, offset);
jacp[i+0*nv] = cdof[3] + tmp[0];
jacp[i+1*nv] = cdof[4] + tmp[1];
jacp[i+2*nv] = cdof[5] + tmp[2];
}
// advance to parent dof
i = m->dof_parentid[i];
}
}
// compute body Jacobian
void mj_jacBody(const mjModel* m, const mjData* d, mjtNum* jacp, mjtNum* jacr, int body) {
mj_jac(m, d, jacp, jacr, d->xpos+3*body, body);
}
// compute body-com Jacobian
void mj_jacBodyCom(const mjModel* m, const mjData* d, mjtNum* jacp, mjtNum* jacr, int body) {
mj_jac(m, d, jacp, jacr, d->xipos+3*body, body);
}
// compute subtree-com Jacobian
void mj_jacSubtreeCom(const mjModel* m, mjData* d, mjtNum* jacp, int body) {
int nv = m->nv;
mj_markStack(d);
mjtNum* jacp_b = mjSTACKALLOC(d, 3*nv, mjtNum);
// clear output
mju_zero(jacp, 3*nv);
// forward pass starting from 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;
}
// b is in the body subtree, add mass-weighted Jacobian into jacp
mj_jac(m, d, jacp_b, NULL, d->xipos+3*b, b);
mju_addToScl(jacp, jacp_b, m->body_mass[b], 3*nv);
}
// normalize by subtree mass
mju_scl(jacp, jacp, 1/m->body_subtreemass[body], 3*nv);
mj_freeStack(d);
}
// compute geom Jacobian
void mj_jacGeom(const mjModel* m, const mjData* d, mjtNum* jacp, mjtNum* jacr, int geom) {
mj_jac(m, d, jacp, jacr, d->geom_xpos + 3*geom, m->geom_bodyid[geom]);
}
// compute site Jacobian
void mj_jacSite(const mjModel* m, const mjData* d, mjtNum* jacp, mjtNum* jacr, int site) {
mj_jac(m, d, jacp, jacr, d->site_xpos + 3*site, m->site_bodyid[site]);
}
// compute translation Jacobian of point, and rotation Jacobian of axis
void mj_jacPointAxis(const mjModel* m, mjData* d, mjtNum* jacPoint, mjtNum* jacAxis,
const mjtNum point[3], const mjtNum axis[3], int body) {
int nv = m->nv;
// get full Jacobian of point
mj_markStack(d);
mjtNum* jacp = (jacPoint ? jacPoint : mjSTACKALLOC(d, 3*nv, mjtNum));
mjtNum* jacr = mjSTACKALLOC(d, 3*nv, mjtNum);
mj_jac(m, d, jacp, jacr, point, body);
// jacAxis_col = cross(jacr_col, axis)
if (jacAxis) {
for (int i=0; i < nv; i++) {
jacAxis[ i] = jacr[ nv+i]*axis[2] - jacr[2*nv+i]*axis[1];
jacAxis[ nv+i] = jacr[2*nv+i]*axis[0] - jacr[ i]*axis[2];
jacAxis[2*nv+i] = jacr[ i]*axis[1] - jacr[ nv+i]*axis[0];
}
}
mj_freeStack(d);
}
// compute 3/6-by-nv sparse Jacobian of global point attached to given body
void mj_jacSparse(const mjModel* m, const mjData* d,
mjtNum* jacp, mjtNum* jacr, const mjtNum* point, int body,
int NV, const int* chain) {
// clear jacobians
if (jacp) {
mju_zero(jacp, 3*NV);
}
if (jacr) {
mju_zero(jacr, 3*NV);
}
// compute point-com offset
mjtNum offset[3];
mju_sub3(offset, point, d->subtree_com+3*m->body_rootid[body]);
// skip fixed bodies
while (body && !m->body_dofnum[body]) {
body = m->body_parentid[body];
}
// no movable body found: nothing to do
if (!body) {
return;
}
// get last dof that affects this (as well as the original) body
int da = m->body_dofadr[body] + m->body_dofnum[body] - 1;
// start and the end of the chain (chain is in increasing order)
int ci = NV-1;
// backward pass over dof ancestor chain
while (da >= 0) {
// find chain index for this dof
while (ci >= 0 && chain[ci] > da) {
ci--;
}
// make sure we found it; SHOULD NOT OCCUR
if (chain[ci] != da) {
mjERROR("dof index %d not found in chain", da);
}
const mjtNum* cdof = d->cdof + 6*da;
// construct rotation jacobian
if (jacr) {
jacr[ci+0*NV] = cdof[0];
jacr[ci+1*NV] = cdof[1];
jacr[ci+2*NV] = cdof[2];
}
// construct translation jacobian (correct for rotation)
if (jacp) {
mjtNum tmp[3];
mju_cross(tmp, cdof, offset);
jacp[ci+0*NV] = cdof[3] + tmp[0];
jacp[ci+1*NV] = cdof[4] + tmp[1];
jacp[ci+2*NV] = cdof[5] + tmp[2];
}
// advance to parent dof
da = m->dof_parentid[da];
}
}
// sparse Jacobian difference for simple body contacts
void mj_jacSparseSimple(const mjModel* m, const mjData* d,
mjtNum* jacdifp, mjtNum* jacdifr, const mjtNum* point,
int body, int flg_second, int NV, int start) {
// compute point-com offset
mjtNum offset[3];
mju_sub3(offset, point, d->subtree_com+3*m->body_rootid[body]);
// skip fixed body
if (!m->body_dofnum[body]) {
return;
}
// process dofs
int ci = start;
int end = m->body_dofadr[body] + m->body_dofnum[body];
for (int da=m->body_dofadr[body]; da < end; da++) {
mjtNum *cdof = d->cdof+6*da;
// construct rotation jacobian
if (jacdifr) {
// plus sign
if (flg_second) {
jacdifr[ci+0*NV] = cdof[0];
jacdifr[ci+1*NV] = cdof[1];
jacdifr[ci+2*NV] = cdof[2];
}
// minus sign
else {
jacdifr[ci+0*NV] = -cdof[0];
jacdifr[ci+1*NV] = -cdof[1];
jacdifr[ci+2*NV] = -cdof[2];
}
}
// construct translation jacobian (correct for rotation)
if (jacdifp) {
mjtNum tmp[3];
mju_cross(tmp, cdof, offset);
// plus sign
if (flg_second) {
jacdifp[ci+0*NV] = (cdof[3] + tmp[0]);
jacdifp[ci+1*NV] = (cdof[4] + tmp[1]);
jacdifp[ci+2*NV] = (cdof[5] + tmp[2]);
}
// plus sign
else {
jacdifp[ci+0*NV] = -(cdof[3] + tmp[0]);
jacdifp[ci+1*NV] = -(cdof[4] + tmp[1]);
jacdifp[ci+2*NV] = -(cdof[5] + tmp[2]);
}
}
// advance jacdif counter
ci++;
}
}
// dense or sparse Jacobian difference for two body points: pos2 - pos1, global
int mj_jacDifPair(const mjModel* m, const mjData* d, int* chain,
int b1, int b2, const mjtNum pos1[3], const mjtNum pos2[3],
mjtNum* jac1p, mjtNum* jac2p, mjtNum* jacdifp,
mjtNum* jac1r, mjtNum* jac2r, mjtNum* jacdifr) {
int issimple = (m->body_simple[b1] && m->body_simple[b2]);
int issparse = mj_isSparse(m);
int NV = m->nv;
// skip if no DOFs
if (!NV) {
return 0;
}
// construct merged chain of body dofs
if (issparse) {
if (issimple) {
NV = mj_mergeChainSimple(m, chain, b1, b2);
} else {
NV = mj_mergeChain(m, chain, b1, b2);
}
}
// skip if empty chain
if (!NV) {
return 0;
}
// sparse case
if (issparse) {
// simple: fast processing
if (issimple) {
// first body
mj_jacSparseSimple(m, d, jacdifp, jacdifr, pos1, b1, 0, NV,
b1 < b2 ? 0 : m->body_dofnum[b2]);
// second body
mj_jacSparseSimple(m, d, jacdifp, jacdifr, pos2, b2, 1, NV,
b2 < b1 ? 0 : m->body_dofnum[b1]);
}
// regular processing
else {
// Jacobians
mj_jacSparse(m, d, jac1p, jac1r, pos1, b1, NV, chain);
mj_jacSparse(m, d, jac2p, jac2r, pos2, b2, NV, chain);
// differences
if (jacdifp) {
mju_sub(jacdifp, jac2p, jac1p, 3*NV);
}
if (jacdifr) {
mju_sub(jacdifr, jac2r, jac1r, 3*NV);
}
}
}
// dense case
else {
// Jacobians
mj_jac(m, d, jac1p, jac1r, pos1, b1);
mj_jac(m, d, jac2p, jac2r, pos2, b2);
// differences
if (jacdifp) {
mju_sub(jacdifp, jac2p, jac1p, 3*NV);
}
if (jacdifr) {
mju_sub(jacdifr, jac2r, jac1r, 3*NV);
}
}
return NV;
}
// dense or sparse weighted sum of multiple body Jacobians at same point
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) {
int nv = m->nv, NV;
mjtNum* jacp = jac;
mjtNum* jacr = flg_rot ? jac + 3*nv : NULL;
mj_markStack(d);
mjtNum* jtmp = mjSTACKALLOC(d, flg_rot ? 6*nv : 3*nv, mjtNum);
mjtNum* jp = jtmp;
mjtNum* jr = flg_rot ? jtmp + 3*nv : NULL;
// sparse
if (mj_isSparse(m)) {
mjtNum* buf = mjSTACKALLOC(d, flg_rot ? 6*nv : 3*nv, mjtNum);
int* buf_ind = mjSTACKALLOC(d, nv, int);
int* bodychain = mjSTACKALLOC(d, nv, int);
// set first
NV = mj_bodyChain(m, body[0], chain);
if (NV) {
// get Jacobian
if (m->body_simple[body[0]]) {
mj_jacSparseSimple(m, d, jacp, jacr, point, body[0], 1, NV, 0);
} else {
mj_jacSparse(m, d, jacp, jacr, point, body[0], NV, chain);
}
// apply weight
mju_scl(jac, jac, weight[0], flg_rot ? 6*NV : 3*NV);
}
// accumulate remaining
for (int i=1; i < n; i++) {
// get body chain and Jacobian
int bodyNV = mj_bodyChain(m, body[i], bodychain);
if (!bodyNV) {
continue;
}
if (m->body_simple[body[i]]) {
mj_jacSparseSimple(m, d, jp, jr, point, body[i], 1, bodyNV, 0);
} else {
mj_jacSparse(m, d, jp, jr, point, body[i], bodyNV, bodychain);
}
// combine sparse matrices
NV = mju_addToSparseMat(jac, jtmp, nv, flg_rot ? 6 : 3, weight[i],
NV, bodyNV, chain, bodychain, buf, buf_ind);
}
}
// dense
else {
// set first
mj_jac(m, d, jacp, jacr, point, body[0]);
mju_scl(jac, jac, weight[0], flg_rot ? 6*nv : 3*nv);
// accumulate remaining
for (int i=1; i < n; i++) {
mj_jac(m, d, jp, jr, point, body[i]);
mju_addToScl(jac, jtmp, weight[i], flg_rot ? 6*nv : 3*nv);
}
NV = nv;
}
mj_freeStack(d);
return NV;
}
// compute 3/6-by-nv Jacobian time derivative of global point attached to given body
void mj_jacDot(const mjModel* m, const mjData* d,
mjtNum* jacp, mjtNum* jacr, const mjtNum point[3], int body) {
int nv = m->nv;
mjtNum offset[3];
mjtNum pvel[6]; // point velocity (rot:lin order)
// clear jacobians, compute offset and pvel if required
if (jacp) {
mju_zero(jacp, 3*nv);
const mjtNum* com = d->subtree_com+3*m->body_rootid[body];
mju_sub3(offset, point, com);
mju_transformSpatial(pvel, d->cvel+6*body, 0, point, com, 0);
}
if (jacr) {
mju_zero(jacr, 3*nv);
}
// skip fixed bodies
while (body && !m->body_dofnum[body]) {
body = m->body_parentid[body];
}
// no movable body found: nothing to do
if (!body) {
return;
}
// get last dof that affects this (as well as the original) body
int i = m->body_dofadr[body] + m->body_dofnum[body] - 1;
// backward pass over dof ancestor chain
while (i >= 0) {
mjtNum cdof_dot[6];
mju_copy(cdof_dot, d->cdof_dot+6*i, 6);
mjtNum* cdof = d->cdof+6*i;
// check for quaternion
mjtJoint type = m->jnt_type[m->dof_jntid[i]];
int dofadr = m->jnt_dofadr[m->dof_jntid[i]];
int is_quat = type == mjJNT_BALL || (type == mjJNT_FREE && i >= dofadr + 3);
// compute cdof_dot for quaternion (use current body cvel)
if (is_quat) {
mju_crossMotion(cdof_dot, d->cvel+6*m->dof_bodyid[i], cdof);
}
// construct rotation jacobian
if (jacr) {
jacr[i+0*nv] += cdof_dot[0];
jacr[i+1*nv] += cdof_dot[1];
jacr[i+2*nv] += cdof_dot[2];
}
// construct translation jacobian (correct for rotation)
if (jacp) {
// first correction term, account for varying cdof
mjtNum tmp1[3];
mju_cross(tmp1, cdof_dot, offset);
// second correction term, account for point translational velocity
mjtNum tmp2[3];
mju_cross(tmp2, cdof, pvel + 3);
jacp[i+0*nv] += cdof_dot[3] + tmp1[0] + tmp2[0];
jacp[i+1*nv] += cdof_dot[4] + tmp1[1] + tmp2[1];
jacp[i+2*nv] += cdof_dot[5] + tmp1[2] + tmp2[2];
}
// advance to parent dof
i = m->dof_parentid[i];
}
}
// 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 = mjSTACKALLOC(d, 3*nv, mjtNum);
mjtNum* jacr = mjSTACKALLOC(d, 3*nv, mjtNum);
mjtNum* term1 = mjSTACKALLOC(d, 3*nv, mjtNum);
mjtNum* term2 = mjSTACKALLOC(d, 3*nv, mjtNum);
// 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_mulMatMat3(tmp1, ximat, inertia); // tmp1 = ximat * inertia
mju_mulMatMatT3(tmp2, tmp1, ximat); // 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);
}
//-------------------------- inertia functions -----------------------------------------------------
// convert sparse inertia matrix M into full matrix
@@ -1118,117 +424,6 @@ void mj_xfrcAccumulate(const mjModel* m, mjData* d, mjtNum* qfrc) {
// compute object 6D velocity in object-centered frame, world/local orientation
void mj_objectVelocity(const mjModel* m, const mjData* d,
int objtype, int objid, mjtNum res[6], int flg_local) {
int bodyid = 0;
const mjtNum *pos = 0, *rot = 0;
// body-inertial
if (objtype == mjOBJ_BODY) {
bodyid = objid;
pos = d->xipos+3*objid;
rot = (flg_local ? d->ximat+9*objid : 0);
}
// body-regular
else if (objtype == mjOBJ_XBODY) {
bodyid = objid;
pos = d->xpos+3*objid;
rot = (flg_local ? d->xmat+9*objid : 0);
}
// geom
else if (objtype == mjOBJ_GEOM) {
bodyid = m->geom_bodyid[objid];
pos = d->geom_xpos+3*objid;
rot = (flg_local ? d->geom_xmat+9*objid : 0);
}
// site
else if (objtype == mjOBJ_SITE) {
bodyid = m->site_bodyid[objid];
pos = d->site_xpos+3*objid;
rot = (flg_local ? d->site_xmat+9*objid : 0);
}
// camera
else if (objtype == mjOBJ_CAMERA) {
bodyid = m->cam_bodyid[objid];
pos = d->cam_xpos+3*objid;
rot = (flg_local ? d->cam_xmat+9*objid : 0);
}
// object without spatial frame
else {
mjERROR("invalid object type %d", objtype);
}
// transform velocity
mju_transformSpatial(res, d->cvel+6*bodyid, 0, pos, d->subtree_com+3*m->body_rootid[bodyid], rot);
}
// compute object 6D acceleration in object-centered frame, world/local orientation
void mj_objectAcceleration(const mjModel* m, const mjData* d,
int objtype, int objid, mjtNum res[6], int flg_local) {
int bodyid = 0;
const mjtNum *pos = 0, *rot = 0;
mjtNum correction[3], vel[6];
// body-inertial
if (objtype == mjOBJ_BODY) {
bodyid = objid;
pos = d->xipos+3*objid;
rot = (flg_local ? d->ximat+9*objid : 0);
}
// body-regular
else if (objtype == mjOBJ_XBODY) {
bodyid = objid;
pos = d->xpos+3*objid;
rot = (flg_local ? d->xmat+9*objid : 0);
}
// geom
else if (objtype == mjOBJ_GEOM) {
bodyid = m->geom_bodyid[objid];
pos = d->geom_xpos+3*objid;
rot = (flg_local ? d->geom_xmat+9*objid : 0);
}
// site
else if (objtype == mjOBJ_SITE) {
bodyid = m->site_bodyid[objid];
pos = d->site_xpos+3*objid;
rot = (flg_local ? d->site_xmat+9*objid : 0);
}
// camera
else if (objtype == mjOBJ_CAMERA) {
bodyid = m->cam_bodyid[objid];
pos = d->cam_xpos+3*objid;
rot = (flg_local ? d->cam_xmat+9*objid : 0);
}
// object without spatial frame
else {
mjERROR("invalid object type %d", objtype);
}
// transform com-based velocity to local frame
mju_transformSpatial(vel, d->cvel+6*bodyid, 0, pos, d->subtree_com+3*m->body_rootid[bodyid], rot);
// transform com-based acceleration to local frame
mju_transformSpatial(res, d->cacc+6*bodyid, 0, pos, d->subtree_com+3*m->body_rootid[bodyid], rot);
// acc_tran += vel_rot x vel_tran
mju_cross(correction, vel, vel+3);
mju_addTo3(res+3, correction);
}
//-------------------------- miscellaneous ---------------------------------------------------------
@@ -1314,28 +509,6 @@ mjtNum mj_geomDistance(const mjModel* m, const mjData* d, int geom1, int geom2,
// extract 6D force:torque for one contact, in contact frame
void mj_contactForce(const mjModel* m, const mjData* d, int id, mjtNum result[6]) {
mjContact* con;
// clear result
mju_zero(result, 6);
// make sure contact is valid
if (id >= 0 && id < d->ncon && d->contact[id].efc_address >= 0) {
// get contact pointer
con = d->contact + id;
if (mj_isPyramidal(m)) {
mju_decodePyramid(result, d->efc_force + con->efc_address, con->friction, con->dim);
} else {
mju_copy(result, d->efc_force + con->efc_address, con->dim);
}
}
}
// compute velocity by finite-differencing two positions
void mj_differentiatePos(const mjModel* m, mjtNum* qvel, mjtNum dt,
const mjtNum* qpos1, const mjtNum* qpos2) {
@@ -1418,50 +591,6 @@ void mj_normalizeQuat(const mjModel* m, mjtNum* qpos) {
// map from body local to global Cartesian coordinates
void mj_local2Global(mjData* d, mjtNum xpos[3], mjtNum xmat[9],
const mjtNum pos[3], const mjtNum quat[4],
int body, mjtByte sameframe) {
mjtSameFrame sf = sameframe;
// position
if (xpos && pos) {
switch (sf) {
case mjSAMEFRAME_NONE:
case mjSAMEFRAME_BODYROT:
case mjSAMEFRAME_INERTIAROT:
mju_mulMatVec3(xpos, d->xmat+9*body, pos);
mju_addTo3(xpos, d->xpos+3*body);
break;
case mjSAMEFRAME_BODY:
mju_copy3(xpos, d->xpos+3*body);
break;
case mjSAMEFRAME_INERTIA:
mju_copy3(xpos, d->xipos+3*body);
break;
}
}
// orientation
if (xmat && quat) {
mjtNum tmp[4];
switch (sf) {
case mjSAMEFRAME_NONE:
mju_mulQuat(tmp, d->xquat+4*body, quat);
mju_quat2Mat(xmat, tmp);
break;
case mjSAMEFRAME_BODY:
case mjSAMEFRAME_BODYROT:
mju_copy(xmat, d->xmat+9*body, 9);
break;
case mjSAMEFRAME_INERTIA:
case mjSAMEFRAME_INERTIAROT:
mju_copy(xmat, d->ximat+9*body, 9);
break;
}
}
}
// return 1 if actuator i is disabled, 0 otherwise
int mj_actuatorDisabled(const mjModel* m, int i) {
int group = m->actuator_group[i];
@@ -1503,27 +632,6 @@ void mj_setTotalmass(mjModel* m, mjtNum newmass) {
// count warnings, print only the first time
void mj_warning(mjData* d, int warning, int info) {
// check type
if (warning < 0 || warning >= mjNWARNING) {
mjERROR("invalid warning type %d", warning);
}
// save info (override previous)
d->warning[warning].lastinfo = info;
// print message only the first time this warning is encountered
if (!d->warning[warning].number) {
mju_warning("%s Time = %.4f.", mju_warningText(warning, info), d->time);
}
// increase counter
d->warning[warning].number++;
}
// version number
int mj_version(void) {
return mjVERSION;
-93
View File
@@ -47,77 +47,6 @@ MJAPI void mj_setState(const mjModel* m, mjData* d, const mjtNum* state, unsigne
// copy current state to the k-th model keyframe
MJAPI void mj_setKeyframe(mjModel* m, const mjData* d, int k);
//-------------------------- sparse chains ---------------------------------------------------------
// merge dof chains for two bodies
int mj_mergeChain(const mjModel* m, int* chain, int b1, int b2);
// merge dof chains for two simple bodies
int mj_mergeChainSimple(const mjModel* m, int* chain, int b1, int b2);
// get body chain
int mj_bodyChain(const mjModel* m, int body, int* chain);
//-------------------------- Jacobians -------------------------------------------------------------
// compute 3/6-by-nv Jacobian of global point attached to given body
MJAPI void mj_jac(const mjModel* m, const mjData* d,
mjtNum* jacp, mjtNum* jacr, const mjtNum point[3], int body);
// compute body frame Jacobian
MJAPI void mj_jacBody(const mjModel* m, const mjData* d,
mjtNum* jacp, mjtNum* jacr, int body);
// compute body center-of-mass Jacobian
MJAPI void mj_jacBodyCom(const mjModel* m, const mjData* d,
mjtNum* jacp, mjtNum* jacr, int body);
// compute subtree center-of-mass Jacobian
MJAPI void mj_jacSubtreeCom(const mjModel* m, mjData* d, mjtNum* jacp, int body);
// compute geom Jacobian
MJAPI void mj_jacGeom(const mjModel* m, const mjData* d,
mjtNum* jacp, mjtNum* jacr, int geom);
// compute site Jacobian
MJAPI void mj_jacSite(const mjModel* m, const mjData* d,
mjtNum* jacp, mjtNum* jacr, int site);
// compute translation Jacobian of point, and rotation Jacobian of axis
MJAPI void mj_jacPointAxis(const mjModel* m, mjData* d,
mjtNum* jacPoint, mjtNum* jacAxis,
const mjtNum point[3], const mjtNum axis[3], int body);
// compute 3/6-by-nv sparse Jacobian of global point attached to given body
void mj_jacSparse(const mjModel* m, const mjData* d,
mjtNum* jacp, mjtNum* jacr, const mjtNum* point, int body,
int NV, const int* chain);
// sparse Jacobian difference for simple body contacts
void mj_jacSparseSimple(const mjModel* m, const mjData* d,
mjtNum* jacdifp, mjtNum* jacdifr, const mjtNum* point,
int body, int flg_second, int NV, int start);
// dense or sparse Jacobian difference for two body points: pos2 - pos1, global
MJAPI int mj_jacDifPair(const mjModel* m, const mjData* d, int* chain,
int b1, int b2, const mjtNum pos1[3], const mjtNum pos2[3],
mjtNum* jac1p, mjtNum* jac2p, mjtNum* jacdifp,
mjtNum* jac1r, mjtNum* jac2r, mjtNum* jacdifr);
// dense or sparse weighted sum of multiple body Jacobians at same point
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 3/6-by-nv Jacobian time derivative of global point attached to given body
MJAPI void mj_jacDot(const mjModel* m, const mjData* d,
mjtNum* jacp, mjtNum* jacr, const mjtNum point[3], int body);
// compute subtree angular momentum matrix
MJAPI void mj_angmomMat(const mjModel* m, mjData* d, mjtNum* mat, int body);
//-------------------------- inertia functions -----------------------------------------------------
// convert sparse inertia matrix M into full matrix
@@ -146,26 +75,12 @@ MJAPI void mj_applyFT(const mjModel* m, mjData* d,
void mj_xfrcAccumulate(const mjModel* m, mjData* d, mjtNum* qfrc);
//-------------------------- coordinate transformation ---------------------------------------------
// compute object 6D velocity in object-centered frame, world/local orientation
MJAPI void mj_objectVelocity(const mjModel* m, const mjData* d,
int objtype, int objid, mjtNum res[6], int flg_local);
// compute object 6D acceleration in object-centered frame, world/local orientation
MJAPI void mj_objectAcceleration(const mjModel* m, const mjData* d,
int objtype, int objid, mjtNum res[6], int flg_local);
//-------------------------- miscellaneous ---------------------------------------------------------
// returns the smallest distance between two geoms
MJAPI mjtNum mj_geomDistance(const mjModel* m, const mjData* d, int geom1, int geom2,
mjtNum distmax, mjtNum fromto[6]);
// extract 6D force:torque for one contact, in contact frame
MJAPI void mj_contactForce(const mjModel* m, const mjData* d, int id, mjtNum result[6]);
// compute velocity by finite-differencing two positions
MJAPI void mj_differentiatePos(const mjModel* m, mjtNum* qvel, mjtNum dt,
const mjtNum* qpos1, const mjtNum* qpos2);
@@ -176,11 +91,6 @@ MJAPI void mj_integratePos(const mjModel* m, mjtNum* qpos, const mjtNum* qvel, m
// normalize all quaternions in qpos-type vector
MJAPI void mj_normalizeQuat(const mjModel* m, mjtNum* qpos);
// map from body local to global Cartesian coordinates
MJAPI void mj_local2Global(mjData* d, mjtNum xpos[3], mjtNum xmat[9],
const mjtNum pos[3], const mjtNum quat[4],
int body, mjtByte sameframe);
// return 1 if actuator i is disabled, 0 otherwise
MJAPI int mj_actuatorDisabled(const mjModel* m, int i);
@@ -190,9 +100,6 @@ MJAPI mjtNum mj_getTotalmass(const mjModel* m);
// scale body masses and inertias to achieve specified total mass
MJAPI void mj_setTotalmass(mjModel* m, mjtNum newmass);
// high-level warning function: count warnings in mjData, print only the first time
MJAPI void mj_warning(mjData* d, int warning, int info);
// version number
MJAPI int mj_version(void);
+1 -1
View File
@@ -20,7 +20,7 @@
#include <mujoco/mujoco.h>
#include "engine/engine_crossplatform.h"
#include "engine/engine_io.h"
#include "engine/engine_memory.h"
// stack allocate and initialize new mjArrayList
mjArrayList* mju_arrayListCreate(mjData* d, size_t element_size, size_t initial_capacity) {
+1 -1
View File
@@ -19,11 +19,11 @@
#include <mujoco/mjdata.h>
#include <mujoco/mjmacro.h>
#include <mujoco/mjsan.h> // IWYU pragma: keep
#include "engine/engine_io.h"
#include "engine/engine_util_blas.h"
#include "engine/engine_util_errmem.h"
#include "engine/engine_util_misc.h"
#include "engine/engine_util_sparse.h"
#include "engine/engine_memory.h"
#include "engine/engine_util_spatial.h"
//---------------------------- dense Cholesky ------------------------------------------------------
+1 -1
View File
@@ -21,7 +21,7 @@
#include <mujoco/mjmacro.h>
#include <mujoco/mjsan.h> // IWYU pragma: keep
#include <mujoco/mjtnum.h>
#include "engine/engine_io.h"
#include "engine/engine_memory.h"
#include "engine/engine_util_blas.h"
#include "engine/engine_util_misc.h"
+2
View File
@@ -22,7 +22,9 @@
#include <mujoco/mjsan.h> // IWYU pragma: keep
#include <mujoco/mjvisualize.h>
#include "engine/engine_core_smooth.h"
#include "engine/engine_core_util.h"
#include "engine/engine_io.h"
#include "engine/engine_memory.h"
#include "engine/engine_ray.h"
#include "engine/engine_support.h"
#include "engine/engine_util_blas.h"
+2 -1
View File
@@ -23,7 +23,8 @@
#include <mujoco/mjsan.h> // IWYU pragma: keep
#include <mujoco/mjvisualize.h>
#include "engine/engine_array_safety.h"
#include "engine/engine_io.h"
#include "engine/engine_core_util.h"
#include "engine/engine_memory.h"
#include "engine/engine_name.h"
#include "engine/engine_plugin.h"
#include "engine/engine_support.h"
+1
View File
@@ -16,6 +16,7 @@
#include <mujoco/mjmodel.h>
#include <mujoco/mjvisualize.h>
#include <mujoco/mjspec.h>
#include "engine/engine_core_util.h"
#include "engine/engine_io.h"
#include "user/user_api.h"
@@ -24,6 +24,7 @@
#include <mujoco/mjmodel.h>
#include <mujoco/mujoco.h>
#include "src/engine/engine_core_constraint.h"
#include "src/engine/engine_core_util.h"
#include "src/engine/engine_support.h"
#include "src/engine/engine_util_misc.h"
#include "test/fixture.h"
@@ -197,7 +198,9 @@ TEST_F(CoreConstraintTest, RestPenetration) {
// simulate for 50 seconds
mj_resetData(model, data);
while (data->time < 50) {
mjtNum time = data->time;
mj_step(model, data);
ASSERT_GT(data->time, time) << "Divergence detected";
}
mjtNum depth = -data->contact[0].dist;
@@ -263,7 +266,9 @@ TEST_F(CoreConstraintTest, EqualityBodySite) {
// simulate, get diag(A)
while (data->time < 0.1) {
mjtNum time = data->time;
mj_step(model, data);
ASSERT_GT(data->time, time) << "Divergence detected";
}
int nefc_site = data->nefc;
std::vector<mjtNum> dA = AsVector(data->efc_diagApprox, nefc_site);
@@ -276,7 +281,9 @@ TEST_F(CoreConstraintTest, EqualityBodySite) {
// simulate again, get diag(A)
while (data->time < 0.1) {
mjtNum time = data->time;
mj_step(model, data);
ASSERT_GT(data->time, time) << "Divergence detected";
}
// compare
@@ -310,8 +317,12 @@ TEST_F(CoreConstraintTest, ConstraintUpdateImpl) {
mj_resetData(model, d1);
mj_resetData(model, d2);
while (d1->time < 0.2) {
mjtNum time1 = d1->time;
mjtNum time2 = d2->time;
mj_step(model, d1);
ASSERT_GT(d1->time, time1) << "Divergence detected";
mj_step(model, d2);
ASSERT_GT(d2->time, time2) << "Divergence detected";
}
mj_forward(model, d1);
mj_forward(model, d2);
+16
View File
@@ -187,7 +187,9 @@ TEST_F(CoreSmoothTest, TendonJdot) {
} else {
mj_resetData(m, d);
while (d->time < 1) {
mjtNum time = d->time;
mj_step(m, d);
ASSERT_GT(d->time, time) << "Divergence detected";
}
}
@@ -307,7 +309,9 @@ TEST_F(CoreSmoothTest, TendonArmatureConservesEnergy) {
double eps = std::max(energy_0, 1.0) * 1e-5;
while (d->time < 1) {
mjtNum time = d->time;
mj_step(m, d);
ASSERT_GT(d->time, time) << "Divergence detected";
double energy_t = d->energy[0] + d->energy[1];
EXPECT_THAT(energy_t, DoubleNear(energy_0, eps));
}
@@ -336,7 +340,9 @@ TEST_F(CoreSmoothTest, TendonArmatureConservesMomentum) {
double eps = 1e-5;
while (d->time < 1) {
mjtNum time = d->time;
mj_step(m, d);
ASSERT_GT(d->time, time) << "Divergence detected";
vector<mjtNum> sdata_t = AsVector(d->sensordata, m->nsensordata);
EXPECT_THAT(sdata_t, Pointwise(DoubleNear(eps), sdata_0));
}
@@ -380,10 +386,14 @@ TEST_F(CoreSmoothTest, TendonInertiaEquivalent) {
double eps = lpath == kTen_i0 ? 1e-6 : 1e-3;
while (d->time < 1) {
mjtNum time = d->time;
mj_step(m, d);
ASSERT_GT(d->time, time) << "Divergence detected";
vector<mjtNum> xpos = AsVector(d->geom_xpos + 3*gid, 3);
time = d_e->time;
mj_step(m_e, d_e);
ASSERT_GT(d_e->time, time) << "Divergence detected";
vector<mjtNum> xpos_e = AsVector(d_e->geom_xpos + 3*gid_e, 3);
EXPECT_THAT(xpos, Pointwise(DoubleNear(eps), xpos_e));
@@ -560,7 +570,9 @@ TEST_F(CoreSmoothTest, EqualityBodySite) {
// simulate again, get sensordata
while (data->time < 0.1) {
mjtNum time = data->time;
mj_step(model, data);
ASSERT_GT(data->time, time) << "Divergence detected";
}
// compare
@@ -589,7 +601,9 @@ TEST_F(CoreSmoothTest, RefsiteBringsToPose) {
// step for 5 seconds
while (data->time < 10) {
mjtNum time = data->time;
mj_step(model, data);
ASSERT_GT(data->time, time) << "Divergence detected";
}
// get site IDs
@@ -628,7 +642,9 @@ TEST_F(CoreSmoothTest, RefsiteConservesMomentum) {
// simulate, assert that momentum is conserved
mjtNum eps = 1e-9;
while (data->time < 1) {
mjtNum time = data->time;
mj_step(model, data);
ASSERT_GT(data->time, time) << "Divergence detected";
for (int i=0; i < 6; i++) {
EXPECT_LT(mju_abs(data->sensordata[i]), eps);
}
+3 -2
View File
@@ -12,15 +12,16 @@
// See the License for the specific language governing permissions and
// limitations under the License.
// Tests for engine/engine_io.c.
// Tests for engine/{engine_io.c and engine_memory.c}.
#include "src/engine/engine_io.h"
#include "src/engine/engine_memory.h"
#include <array>
#include <cstdint>
#include <cstdio>
#include <cstring>
#include <filesystem>
#include <filesystem> // NOLINT
#include <string>
#include <vector>
+2 -1
View File
@@ -12,8 +12,9 @@
// See the License for the specific language governing permissions and
// limitations under the License.
// Tests for engine/engine_support.c.
// Tests for engine/{engine_support.c and engine_core_util.c}
#include "src/engine/engine_core_util.h"
#include "src/engine/engine_support.h"
#include <limits>