Move elasticity force computation to the engine.

PiperOrigin-RevId: 676028775
Change-Id: I994e3dda8d6470815763c04df0c276e265aa1f54
This commit is contained in:
Alessio Quaglino
2024-09-18 09:53:11 -07:00
committed by Copybara-Service
parent e439b748a4
commit 66a9cacc90
6 changed files with 127 additions and 38 deletions
+7 -13
View File
@@ -31,15 +31,10 @@ namespace mujoco::plugin::elasticity {
// factory function
std::optional<Membrane> Membrane::Create(const mjModel* m, mjData* d,
int instance) {
if (CheckAttr("face", m, instance) && CheckAttr("poisson", m, instance) &&
CheckAttr("young", m, instance) && CheckAttr("thickness", m, instance)) {
mjtNum nu = strtod(mj_getPluginConfig(m, instance, "poisson"), nullptr);
mjtNum E = strtod(mj_getPluginConfig(m, instance, "young"), nullptr);
mjtNum thick =
strtod(mj_getPluginConfig(m, instance, "thickness"), nullptr);
if (CheckAttr("face", m, instance)) {
mjtNum damp =
strtod(mj_getPluginConfig(m, instance, "damping"), nullptr);
return Membrane(m, d, instance, nu, E, thick, damp);
return Membrane(m, d, instance, damp);
} else {
mju_warning("Invalid parameter specification in shell plugin");
return std::nullopt;
@@ -47,9 +42,8 @@ std::optional<Membrane> Membrane::Create(const mjModel* m, mjData* d,
}
// plugin constructor
Membrane::Membrane(const mjModel* m, mjData* d, int instance, mjtNum nu,
mjtNum E, mjtNum thick, mjtNum damp)
: f0(-1), damping(damp), thickness(thick) {
Membrane::Membrane(const mjModel* m, mjData* d, int instance, mjtNum damp)
: f0(-1), damping(damp) {
// count plugin bodies
nv = ne = 0;
for (int i = 1; i < m->nbody; i++) {
@@ -106,8 +100,8 @@ void Membrane::Compute(const mjModel* m, mjData* d, int instance) {
// Animation" http://multires.caltech.edu/pubs/DiscreteLagrangian.pdf
for (int idx = 0; idx < ne; idx++) {
elongation[idx] = deformed[idx]*deformed[idx] - ref[idx]*ref[idx] +
( deformed[idx]*deformed[idx] - prev[idx]*prev[idx] ) * kD;
elongation[idx] =
(deformed[idx] * deformed[idx] - prev[idx] * prev[idx]) * kD;
}
// compute gradient of elastic energy and insert into passive force
@@ -135,7 +129,7 @@ void Membrane::RegisterPlugin() {
plugin.name = "mujoco.elasticity.membrane";
plugin.capabilityflags |= mjPLUGIN_PASSIVE;
const char* attributes[] = {"face", "edge", "young", "poisson", "thickness", "damping"};
const char* attributes[] = {"face", "edge", "damping"};
plugin.nattribute = sizeof(attributes) / sizeof(attributes[0]);
plugin.attributes = attributes;
plugin.nstate = +[](const mjModel* m, int instance) { return 0; };
+1 -3
View File
@@ -52,11 +52,9 @@ class Membrane {
std::vector<mjtNum> force; // force at all vertices (nv x 3)
mjtNum damping;
mjtNum thickness;
private:
Membrane(const mjModel* m, mjData* d, int instance, mjtNum nu, mjtNum E,
mjtNum thick, mjtNum damp);
Membrane(const mjModel* m, mjData* d, int instance, mjtNum damp);
};
} // namespace mujoco::plugin::elasticity
+6 -11
View File
@@ -32,14 +32,10 @@ namespace mujoco::plugin::elasticity {
// factory function
std::optional<Solid> Solid::Create(const mjModel* m, mjData* d, int instance) {
if (CheckAttr("face", m, instance) &&
CheckAttr("edge", m, instance) &&
CheckAttr("poisson", m, instance) &&
CheckAttr("young", m, instance)) {
mjtNum nu = strtod(mj_getPluginConfig(m, instance, "poisson"), nullptr);
mjtNum E = strtod(mj_getPluginConfig(m, instance, "young"), nullptr);
CheckAttr("edge", m, instance)) {
mjtNum damp =
strtod(mj_getPluginConfig(m, instance, "damping"), nullptr);
return Solid(m, d, instance, nu, E, damp);
return Solid(m, d, instance, damp);
} else {
mju_warning("Invalid parameter specification in solid plugin");
return std::nullopt;
@@ -47,8 +43,7 @@ std::optional<Solid> Solid::Create(const mjModel* m, mjData* d, int instance) {
}
// plugin constructor
Solid::Solid(const mjModel* m, mjData* d, int instance, mjtNum nu, mjtNum E,
mjtNum damp)
Solid::Solid(const mjModel* m, mjData* d, int instance, mjtNum damp)
: f0(-1), damping(damp) {
// count plugin bodies
nv = ne = 0;
@@ -109,8 +104,8 @@ void Solid::Compute(const mjModel* m, mjData* d, int instance) {
// Animation" http://multires.caltech.edu/pubs/DiscreteLagrangian.pdf
for (int idx = 0; idx < ne; idx++) {
elongation[idx] = deformed[idx]*deformed[idx] - ref[idx]*ref[idx] +
( deformed[idx]*deformed[idx] - prev[idx]*prev[idx] ) * kD;
elongation[idx] =
(deformed[idx] * deformed[idx] - prev[idx] * prev[idx]) * kD;
}
// compute gradient of elastic energy and insert into passive force
@@ -138,7 +133,7 @@ void Solid::RegisterPlugin() {
plugin.name = "mujoco.elasticity.solid";
plugin.capabilityflags |= mjPLUGIN_PASSIVE;
const char* attributes[] = {"face", "edge", "young", "poisson", "damping", "thickness"};
const char* attributes[] = {"face", "edge", "damping"};
plugin.nattribute = sizeof(attributes) / sizeof(attributes[0]);
plugin.attributes = attributes;
plugin.nstate = +[](const mjModel* m, int instance) { return 0; };
+1 -2
View File
@@ -52,8 +52,7 @@ class Solid {
mjtNum damping;
private:
Solid(const mjModel* m, mjData* d, int instance, mjtNum nu, mjtNum E,
mjtNum damp);
Solid(const mjModel* m, mjData* d, int instance, mjtNum damp);
};
} // namespace mujoco::plugin::elasticity
+112
View File
@@ -23,6 +23,7 @@
#include "engine/engine_callback.h"
#include "engine/engine_core_constraint.h"
#include "engine/engine_crossplatform.h"
#include "engine/engine_io.h"
#include "engine/engine_plugin.h"
#include "engine/engine_support.h"
#include "engine/engine_util_blas.h"
@@ -33,6 +34,25 @@
//----------------------------- passive forces -----------------------------------------------------
// local edge-based vertex indexing for 2D and 3D elements, 2D and 3D elements
// have 3 and 6 edges, respectively so the missing indexes are set to 0
static int edges[2][6][2] = {{{1, 2}, {2, 0}, {0, 1}, {0, 0}, {0, 0}, {0, 0}},
{{0, 1}, {1, 2}, {2, 0}, {2, 3}, {0, 3}, {1, 3}}};
// compute gradient of squared lengths of edges belonging to a given element
static void inline GradSquaredLengths(mjtNum gradient[6][2][3],
const mjtNum* xpos,
const int vert[4],
const int edge[6][2],
int nedge) {
for (int e = 0; e < nedge; e++) {
for (int d = 0; d < 3; d++) {
gradient[e][0][d] = xpos[3*vert[edge[e][0]]+d] - xpos[3*vert[edge[e][1]]+d];
gradient[e][1][d] = xpos[3*vert[edge[e][1]]+d] - xpos[3*vert[edge[e][0]]+d];
}
}
}
// spring and damper forces
static void mj_springdamper(const mjModel* m, mjData* d) {
int nv = m->nv, njnt = m->njnt, ntendon = m->ntendon;
@@ -93,6 +113,98 @@ static void mj_springdamper(const mjModel* m, mjData* d) {
}
}
// flex elasticity
for (int f=0; f < m->nflex; f++) {
mjtNum* k = m->flex_stiffness + 21*m->flex_elemadr[f];
int dim = m->flex_dim[f];
if (dim == 1 || m->flex_rigid[f] || k[0] == 0) {
continue;
}
int nedge = (dim == 2) ? 3 : 6;
int nvert = (dim == 2) ? 3 : 4;
const int* elem = m->flex_elem + m->flex_elemdataadr[f];
const int* edgeelem = m->flex_elemedge + m->flex_elemedgeadr[f];
mjtNum* xpos = d->flexvert_xpos + 3*m->flex_vertadr[f];
mjtNum* deformed = d->flexedge_length + m->flex_edgeadr[f];
mjtNum* ref = m->flexedge_length0 + m->flex_edgeadr[f];
int* bodyid = m->flex_vertbodyid + m->flex_vertadr[f];
mj_markStack(d);
mjtNum* qfrc = mj_stackAllocNum(d, 3*m->flex_vertnum[f]);
mju_zero(qfrc, 3*m->flex_vertnum[f]);
// compute force element-by-element
for (int t = 0; t < m->flex_elemnum[f]; t++) {
const int* vert = elem + (dim+1) * t;
// compute length gradient with respect to dofs
mjtNum gradient[6][2][3];
GradSquaredLengths(gradient, xpos, vert, edges[dim-2], nedge);
// extract elongation of edges belonging to this element
mjtNum elongation[6];
for (int e = 0; e < nedge; e++) {
int idx = edgeelem[t * nedge + e];
elongation[e] = deformed[idx]*deformed[idx] - ref[idx]*ref[idx];
}
// unpack triangular representation
mjtNum metric[36];
int id = 0;
for (int ed1 = 0; ed1 < nedge; ed1++) {
for (int ed2 = ed1; ed2 < nedge; ed2++) {
metric[nedge*ed1 + ed2] = k[21*t + id];
metric[nedge*ed2 + ed1] = k[21*t + id++];
}
}
// we now multiply the elongations by the precomputed metric tensor,
// notice that if metric=diag(1/reference) then this would yield a
// mass-spring model
// compute local force
mjtNum force[12] = {0};
for (int ed1 = 0; ed1 < nedge; ed1++) {
for (int ed2 = 0; ed2 < nedge; ed2++) {
for (int i = 0; i < 2; i++) {
for (int x = 0; x < 3; x++) {
force[3 * edges[dim-2][ed2][i] + x] -=
elongation[ed1] * gradient[ed2][i][x] *
metric[nedge * ed1 + ed2];
}
}
}
}
// insert into global force
for (int i = 0; i < nvert; i++) {
for (int x = 0; x < 3; x++) {
qfrc[3*vert[i]+x] += force[3*i+x];
}
}
}
// insert force into qfrc_passive, straightforward for simple bodies,
// need to distribute the force in case of pinned vertices
for (int v = 0; v < m->flex_vertnum[f]; v++) {
int bid = bodyid[v];
if (m->body_simple[bid] != 2) {
// this should only occur for pinned flex vertices
mj_applyFT(m, d, qfrc + 3*v, 0, xpos + 3*v, bid, d->qfrc_spring);
} else {
int body_dofnum = m->body_dofnum[bid];
int body_dofadr = m->body_dofadr[bid];
for (int x = 0; x < body_dofnum; x++) {
d->qfrc_spring[body_dofadr+x] += qfrc[3*v+x];
}
}
}
mj_freeStack(d);
}
// flexedge-level spring-dampers
for (int f=0; f < m->nflex; f++) {
mjtNum stiffness = m->flex_edgestiffness[f];
-9
View File
@@ -2898,15 +2898,6 @@ void mjCFlex::Compile(const mjVFS* vfs) {
if (model->Bodies()[vbodyid]->plugin.instance) {
mjCPlugin* plugin_instance =
static_cast<mjCPlugin*>(model->Bodies()[vbodyid]->plugin.instance);
if (young > 0) {
plugin_instance->config_attribs["young"] = std::to_string(young);
}
if (poisson > 0) {
plugin_instance->config_attribs["poisson"] = std::to_string(poisson);
}
if (thickness > 0) {
plugin_instance->config_attribs["thickness"] = std::to_string(thickness);
}
if (damping > 0) {
plugin_instance->config_attribs["damping"] = std::to_string(damping);
}