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

Fix flex constraint Jacobian for edges belonging to a single body.

Browse Source 
PiperOrigin-RevId: 575161094
Change-Id: If4153bcb751cdd6dbdddeb08672f7107c7c5bbc2
This commit is contained in:
Alessio Quaglino
2023-10-20 02:48:46 -07:00
committed by Copybara-Service
parent e4a9f535a4
commit 01a932e052
12 changed files with 97 additions and 38 deletions
Download Patch File Download Diff File Expand all files Collapse all files
doc/includes/references.h
+2 -1
View File
@@ -944,7 +944,7 @@ struct mjModel_ {
int* body_treeid; // id of body's kinematic tree; -1: static (nbody x 1)
int* body_geomnum; // number of geoms (nbody x 1)
int* body_geomadr; // start addr of geoms; -1: no geoms (nbody x 1)
mjtByte* body_simple; // 1: diagonal M; 2: diag M, no rotations (nbody x 1)
mjtByte* body_simple; // 1: diag M; 2: diag M, sliders only (nbody x 1)
mjtByte* body_sameframe; // inertial frame is same as body frame (nbody x 1)
mjtNum* body_pos; // position offset rel. to parent body (nbody x 3)
mjtNum* body_quat; // orientation offset rel. to parent body (nbody x 4)
@@ -1122,6 +1122,7 @@ struct mjModel_ {
mjtNum* flex_edgedamping; // edge damping (nflex x 1)
mjtByte* flex_edgeequality; // is edge equality constraint defined (nflex x 1)
mjtByte* flex_rigid; // are all verices in the same body (nflex x 1)
mjtByte* flexedge_rigid; // are both edge vertices in same body (nflexedge x 1)
mjtByte* flex_centered; // are all vertex coordinates (0,0,0) (nflex x 1)
mjtByte* flex_flatskin; // render flex skin with flat shading (nflex x 1)
int* flex_bvhadr; // address of bvh root; -1: no bvh (nflex x 1)
include/mujoco/mjmodel.h
+2 -1
View File
@@ -661,7 +661,7 @@ struct mjModel_ {
int* body_treeid; // id of body's kinematic tree; -1: static (nbody x 1)
int* body_geomnum; // number of geoms (nbody x 1)
int* body_geomadr; // start addr of geoms; -1: no geoms (nbody x 1)
mjtByte* body_simple; // 1: diagonal M; 2: diag M, no rotations (nbody x 1)
mjtByte* body_simple; // 1: diag M; 2: diag M, sliders only (nbody x 1)
mjtByte* body_sameframe; // inertial frame is same as body frame (nbody x 1)
mjtNum* body_pos; // position offset rel. to parent body (nbody x 3)
mjtNum* body_quat; // orientation offset rel. to parent body (nbody x 4)
@@ -839,6 +839,7 @@ struct mjModel_ {
mjtNum* flex_edgedamping; // edge damping (nflex x 1)
mjtByte* flex_edgeequality; // is edge equality constraint defined (nflex x 1)
mjtByte* flex_rigid; // are all verices in the same body (nflex x 1)
mjtByte* flexedge_rigid; // are both edge vertices in same body (nflexedge x 1)
mjtByte* flex_centered; // are all vertex coordinates (0,0,0) (nflex x 1)
mjtByte* flex_flatskin; // render flex skin with flat shading (nflex x 1)
int* flex_bvhadr; // address of bvh root; -1: no bvh (nflex x 1)
include/mujoco/mjxmacro.h
+1
View File
@@ -345,6 +345,7 @@
X ( mjtNum, flex_edgedamping, nflex, 1 ) \
X ( mjtByte, flex_edgeequality, nflex, 1 ) \
X ( mjtByte, flex_rigid, nflex, 1 ) \
X ( mjtByte, flexedge_rigid, nflexedge, 1 ) \
X ( mjtByte, flex_centered, nflex, 1 ) \
XMJV( mjtByte, flex_flatskin, nflex, 1 ) \
XMJV( int, flex_bvhadr, nflex, 1 ) \
introspect/structs.py
+8 -1
View File
@@ -1311,7 +1311,7 @@ STRUCTS: Mapping[str, StructDecl] = dict([
type=PointerType(
inner_type=ValueType(name='mjtByte'),
),
doc='1: diagonal M; 2: diag M, no rotations (nbody x 1)',
doc='1: diag M; 2: diag M, sliders only (nbody x 1)',
),
StructFieldDecl(
name='body_sameframe',
@@ -2426,6 +2426,13 @@ STRUCTS: Mapping[str, StructDecl] = dict([
),
doc='are all verices in the same body (nflex x 1)',
),
StructFieldDecl(
name='flexedge_rigid',
type=PointerType(
inner_type=ValueType(name='mjtByte'),
),
doc='are both edge vertices in same body (nflexedge x 1)',
),
StructFieldDecl(
name='flex_centered',
type=PointerType(
src/engine/engine_core_constraint.c
+33 -19
View File
@@ -663,9 +663,13 @@ void mj_instantiateEquality(const mjModel* m, mjData* d) {
case mjEQ_FLEX:
flex_edgeadr = m->flex_edgeadr[id[0]];
flex_edgenum = m->flex_edgenum[id[0]];
// add one constraint per edge
// add one constraint per non-rigid edge
for (int e=flex_edgeadr; e < flex_edgeadr+flex_edgenum; e++) {
// skip rigid
if (m->flexedge_rigid[e]) {
continue;
}
// position error
cpos[0] = d->flexedge_length[e] - m->flexedge_length0[e];
@@ -1057,7 +1061,7 @@ void mj_instantiateContact(const mjModel* m, mjData* d) {
// compute diagApprox
void mj_diagApprox(const mjModel* m, mjData* d) {
int id, dim, b1, b2, weldcnt = 0, edgecnt = 0;
int id, dim, b1, b2, f, weldcnt = 0;
int nefc = d->nefc;
mjtNum tran, rot, fri, *dA = d->efc_diagApprox;
mjContact* con = NULL;
@@ -1067,11 +1071,6 @@ void mj_diagApprox(const mjModel* m, mjData* d) {
// get constraint id
id = d->efc_id[i];
// clear edge counter
if (d->efc_type[i] != mjEQ_FLEX) {
edgecnt = 0;
}
// process according to constraint type
switch ((mjtConstraint) d->efc_type[i]) {
case mjCNSTR_EQUALITY:
@@ -1108,8 +1107,18 @@ void mj_diagApprox(const mjModel* m, mjData* d) {
break;
case mjEQ_FLEX:
dA[i] = m->flexedge_invweight0[m->flex_edgeadr[m->eq_obj1id[id]] + edgecnt];
edgecnt++;
// process all non-rigid edges for this flex
f = m->eq_obj1id[id];
int flex_edgeadr = m->flex_edgeadr[f];
int flex_edgenum = m->flex_edgenum[f];
for (int e=flex_edgeadr; e<flex_edgeadr+flex_edgenum; e++) {
if (!m->flexedge_rigid[e]) {
dA[i++] = m->flexedge_invweight0[e];
}
}
// adjust constraint counter
i--;
break;
default:
@@ -1642,21 +1651,26 @@ static int mj_ne(const mjModel* m, mjData* d, int* nnz) {
break;
case mjEQ_FLEX:
size = m->flex_edgenum[id[0]];
if (!nnz) {
break;
}
flex_edgeadr = m->flex_edgeadr[id[0]];
flex_edgenum = m->flex_edgenum[id[0]];
// init with all edges, subract rigid later
size = flex_edgenum;
// process edges of this flex
for (int e=flex_edgeadr; e < flex_edgeadr+flex_edgenum; e++) {
int b1 = m->flex_vertbodyid[m->flex_vertadr[id[0]] + m->flex_edge[2*e]];
int b2 = m->flex_vertbodyid[m->flex_vertadr[id[0]] + m->flex_edge[2*e+1]];
// rigid: reduce size and skip
if (m->flexedge_rigid[e]) {
size--;
continue;
}
// accumulate NV
NV += mj_jacDifPairCount(m, chain, b1, b2, issparse);
// accumulate NV if needed
if (nnz) {
int b1 = m->flex_vertbodyid[m->flex_vertadr[id[0]] + m->flex_edge[2*e]];
int b2 = m->flex_vertbodyid[m->flex_vertadr[id[0]] + m->flex_edge[2*e+1]];
NV += mj_jacDifPairCount(m, chain, b1, b2, issparse);
}
}
break;
src/engine/engine_core_smooth.c
+10 -2
View File
@@ -1857,8 +1857,16 @@ void mj_rnePostConstraint(const mjModel* m, mjData* d) {
break;
case mjEQ_FLEX:
// increment edgenum rows
i += m->flex_edgenum[m->eq_obj1id[id]];
// increment with number of non-rigid edges
k = m->eq_obj1id[id];
int flex_edgeadr = m->flex_edgeadr[k];
int flex_edgenum = m->flex_edgenum[k];
for (int e=flex_edgeadr; e < flex_edgeadr+flex_edgenum; e++) {
if (!m->flexedge_rigid[e]) {
i++;
}
}
break;
default:
src/engine/engine_derivative.c
+9 -2
View File
@@ -1474,9 +1474,16 @@ void mjd_passive_vel(const mjModel* m, mjData* d) {
for (int f=0; f < m->nflex; f++) {
if (!m->flex_rigid[f] && m->flex_edgedamping[f]) {
mjtNum B = -m->flex_edgedamping[f];
int flex_edgeadr = m->flex_edgeadr[f];
int flex_edgenum = m->flex_edgenum[f];
// process non-rigid edges of this flex
for (int e=flex_edgeadr; e < flex_edgeadr+flex_edgenum; e++) {
// skip rigid
if (m->flexedge_rigid[e]) {
continue;
}
// process edges of this flex
for (int e=m->flex_edgeadr[f]; e < m->flex_edgeadr[f]+m->flex_edgenum[f]; e++) {
// add sparse or dense
if (mj_isSparse(m)) {
addJTBJSparse(m, d, d->flexedge_J, &B, 1, e,
src/engine/engine_passive.c
+6 -1
View File
@@ -106,9 +106,14 @@ void mj_passive(const mjModel* m, mjData* d) {
continue;
}
// process edges of this flex (global edge index)
// process non-rigid edges of this flex (global edge index)
int edgeend = m->flex_edgeadr[f] + m->flex_edgenum[f];
for (int e=m->flex_edgeadr[f]; e < edgeend; e++) {
// skip rigid
if (m->flexedge_rigid[e]) {
continue;
}
// compute spring-damper force along edge
frc = stiffness * (m->flexedge_length0[e] - d->flexedge_length[e])
- damping * d->flexedge_velocity[e];
src/engine/engine_sensor.c
+8 -4
View File
@@ -956,10 +956,14 @@ void mj_energyPos(const mjModel* m, mjData* d) {
continue;
}
// process edges of this flex
for (int e=m->flex_edgeadr[i]; e < m->flex_edgeadr[i]+m->flex_edgenum[i]; e++) {
mjtNum displacement = m->flexedge_length0[e] - d->flexedge_length[e];
d->energy[0] += 0.5*stiffness*displacement*displacement;
// process non-rigid edges of this flex
int flex_edgeadr = m->flex_edgeadr[i];
int flex_edgenum = m->flex_edgenum[i];
for (int e=flex_edgeadr; e < flex_edgeadr+flex_edgenum; e++) {
if (!m->flexedge_rigid[e]) {
mjtNum displacement = m->flexedge_length0[e] - d->flexedge_length[e];
d->energy[0] += 0.5*stiffness*displacement*displacement;
};
}
}
}
src/engine/engine_setconst.c
+6 -1
View File
@@ -206,8 +206,13 @@ static void set0(mjModel* m, mjData* d) {
int b1 = m->flex_vertbodyid[m->flex_vertadr[f] + m->flex_edge[2*i]];
int b2 = m->flex_vertbodyid[m->flex_vertadr[f] + m->flex_edge[2*i+1]];
// rigid edge: set to 0
if (m->flexedge_rigid[i]) {
m->flexedge_invweight0[i] = 0;
}
// accelerate edges that connect simple bodies with no rotations
if (m->body_simple[b1] == 2 && m->body_simple[b2] == 2) {
else if (m->body_simple[b1] == 2 && m->body_simple[b2] == 2) {
m->flexedge_invweight0[i] = (1/m->body_mass[b1] + 1/m->body_mass[b2])/2;
}
src/user/user_model.cc
+11 -6
View File
@@ -1569,8 +1569,8 @@ void mjCModel::CopyTree(mjModel* m) {
qposadr += nPOS[pj->type];
}
// simple body with no rotational dofs: promote to simple level 2
if (m->body_simple[i]) {
// simple body with sliders and no rotational dofs: promote to simple level 2
if (m->body_simple[i] && m->body_dofnum[i]) {
m->body_simple[i] = 2;
for (int j=0; j<(int)pb->joints.size(); j++) {
if (pb->joints[j]->type!=mjJNT_SLIDE) {
@@ -1979,10 +1979,15 @@ void mjCModel::CopyObjects(mjModel* m) {
memcpy(m->flex_vertbodyid + vert_adr, pfl->vertbodyid.data(), pfl->nvert*sizeof(int));
}
// convert edge pairs to int array
for (int i=0; i<pfl->nedge; i++) {
m->flex_edge[2*(edge_adr+i)] = pfl->edge[i].first;
m->flex_edge[2*(edge_adr+i)+1] = pfl->edge[i].second;
// convert edge pairs to int array, set edge rigid
for (int k=0; k<pfl->nedge; k++) {
m->flex_edge[2*(edge_adr+k)] = pfl->edge[k].first;
m->flex_edge[2*(edge_adr+k)+1] = pfl->edge[k].second;
// check if vertex body weldids are the same
int b1 = pfl->vertbodyid[pfl->edge[k].first];
int b2 = pfl->vertbodyid[pfl->edge[k].second];
m->flexedge_rigid[edge_adr+k] = (bodies[b1]->weldid == bodies[b2]->weldid);
}
// advance counters
unity/Runtime/Bindings/MjBindings.cs
+1
View File
@@ -5337,6 +5337,7 @@ public unsafe struct mjModel_ {
public double* flex_edgedamping;
public byte* flex_edgeequality;
public byte* flex_rigid;
public byte* flexedge_rigid;
public byte* flex_centered;
public byte* flex_flatskin;
public int* flex_bvhadr;