Parallelize tactile sensor computation.
The computation of tactile sensor values for each taxel is now parallelized using the MuJoCo threadpool when the number of taxels exceeds a threshold. Each thread processes a batch of taxels, accumulating forces from all colliding geoms. PiperOrigin-RevId: 888606457 Change-Id: I98abf3e98a7318fc14080054bea7093ab46e8b8b
This commit is contained in:
Alessio Quaglino
committed by
Copybara-Service
Copybara-Service
parent
f4e3469239
commit
2d33b50243
+187
-114
@@ -35,6 +35,8 @@
|
||||
#include "engine/engine_util_errmem.h"
|
||||
#include "engine/engine_util_misc.h"
|
||||
#include "engine/engine_util_spatial.h"
|
||||
#include "thread/thread_pool.h"
|
||||
#include "thread/thread_task.h"
|
||||
|
||||
|
||||
|
||||
@@ -60,6 +62,143 @@ static int ContactInfoCompare(const ContactInfo* a, const ContactInfo* b, void*
|
||||
mjPARTIAL_SORT(ContactSelect, ContactInfo, ContactInfoCompare);
|
||||
|
||||
|
||||
// arguments for parallel tactile sensor computation
|
||||
typedef struct mjTactileTaskArgs_ {
|
||||
const mjModel* m;
|
||||
mjData* d;
|
||||
int sensor_id;
|
||||
int mesh_id;
|
||||
int geom_id;
|
||||
int parent_weld;
|
||||
int ncontact;
|
||||
int nchannel;
|
||||
int* contact_geom_ids;
|
||||
int start_taxel;
|
||||
int end_taxel;
|
||||
mjtNum* forcesT;
|
||||
} mjTactileTaskArgs;
|
||||
|
||||
|
||||
// worker function for parallel tactile computation over taxel batches
|
||||
static void* tactile_taxel_batch(void* args) {
|
||||
mjTactileTaskArgs* t = (mjTactileTaskArgs*)args;
|
||||
const mjModel* m = t->m;
|
||||
mjData* d = t->d;
|
||||
int mesh_id = t->mesh_id;
|
||||
int geom_id = t->geom_id;
|
||||
int parent_weld = t->parent_weld;
|
||||
int ncon = m->mesh_vertnum[mesh_id];
|
||||
|
||||
mjtNum* geom_pos = d->geom_xpos + 3*geom_id;
|
||||
mjtNum* geom_mat = d->geom_xmat + 9*geom_id;
|
||||
float* mesh_vert = m->mesh_vert + 3*m->mesh_vertadr[mesh_id];
|
||||
float* mesh_normal = m->mesh_normal + 3*m->mesh_normaladr[mesh_id];
|
||||
|
||||
int has_frame = (m->mesh_normalnum[mesh_id] == 3 * m->mesh_vertnum[mesh_id]);
|
||||
int normal_stride = has_frame ? 9 : 3;
|
||||
|
||||
// process taxels in [start_taxel, end_taxel)
|
||||
for (int j = t->start_taxel; j < t->end_taxel; j++) {
|
||||
mjtNum pos[3] = {mesh_vert[3*j + 0], mesh_vert[3*j + 1], mesh_vert[3*j + 2]};
|
||||
|
||||
mjtNum xpos[3];
|
||||
mju_mulMatVec3(xpos, geom_mat, pos);
|
||||
mju_addTo3(xpos, geom_pos);
|
||||
|
||||
// iterate over colliding geoms
|
||||
for (int g = 0; g < t->ncontact; g++) {
|
||||
int geom = t->contact_geom_ids[g];
|
||||
int body = m->geom_bodyid[geom];
|
||||
|
||||
// set up SDF for this contact geom
|
||||
int sdf_instance[2] = {-1, -1};
|
||||
mjtGeom geomtype[2] = {mjGEOM_SDF, mjGEOM_SPHERE};
|
||||
const mjpPlugin* sdf_ptr[2] = {NULL, NULL};
|
||||
if (m->geom_type[geom] == mjGEOM_SDF) {
|
||||
sdf_instance[0] = m->geom_plugin[geom];
|
||||
sdf_ptr[0] = mjc_getSDF(m, geom);
|
||||
} else if (m->geom_type[geom] == mjGEOM_MESH) {
|
||||
sdf_instance[0] = m->geom_dataid[geom];
|
||||
geomtype[0] = (mjtGeom)m->geom_type[geom];
|
||||
} else {
|
||||
sdf_instance[0] = geom;
|
||||
geomtype[0] = (mjtGeom)m->geom_type[geom];
|
||||
}
|
||||
|
||||
// skip mesh geoms not having an octree
|
||||
if (geomtype[0] == mjGEOM_MESH &&
|
||||
m->mesh_octadr[m->geom_dataid[geom]] == -1) {
|
||||
continue;
|
||||
}
|
||||
|
||||
mjSDF geom_sdf;
|
||||
geom_sdf.id = &sdf_instance[0];
|
||||
geom_sdf.type = mjSDFTYPE_SINGLE;
|
||||
geom_sdf.plugin = &sdf_ptr[0];
|
||||
geom_sdf.geomtype = &geomtype[0];
|
||||
|
||||
// position in other geom frame
|
||||
mjtNum tmp[3], lpos[3];
|
||||
mju_sub3(tmp, xpos, d->geom_xpos + 3*geom);
|
||||
mju_mulMatTVec3(lpos, d->geom_xmat + 9*geom, tmp);
|
||||
|
||||
// SDF plugins are in the original mesh frame
|
||||
if (sdf_ptr[0] != NULL) {
|
||||
mjtNum mesh_mat[9];
|
||||
mju_quat2Mat(mesh_mat, m->mesh_quat + 4 * m->geom_dataid[geom]);
|
||||
mju_mulMatVec3(lpos, mesh_mat, lpos);
|
||||
mju_addTo3(lpos, m->mesh_pos + 3 * m->geom_dataid[geom]);
|
||||
}
|
||||
|
||||
// compute distance
|
||||
mjtNum depth = mju_min(mjc_distance(m, d, &geom_sdf, lpos), 0);
|
||||
if (depth == 0) {
|
||||
continue;
|
||||
}
|
||||
|
||||
// get velocity in global frame
|
||||
mjtNum vel_sensor[6], vel_other[6], vel_rel[3];
|
||||
mju_transformSpatial(
|
||||
vel_sensor, d->cvel + 6 * parent_weld, 0, xpos,
|
||||
d->subtree_com + 3 * m->body_rootid[parent_weld], NULL);
|
||||
mju_transformSpatial(
|
||||
vel_other, d->cvel + 6 * body, 0, d->geom_xpos + 3 * geom,
|
||||
d->subtree_com + 3 * m->body_rootid[body], NULL);
|
||||
mju_sub3(vel_rel, vel_sensor+3, vel_other+3);
|
||||
|
||||
// get normal
|
||||
mjtNum normal[3] = {mesh_normal[normal_stride*j + 0],
|
||||
mesh_normal[normal_stride*j + 1],
|
||||
mesh_normal[normal_stride*j + 2]};
|
||||
mju_rotVecQuat(normal, normal, m->mesh_quat + 4 * mesh_id);
|
||||
|
||||
// get contact force
|
||||
mjtNum force[3];
|
||||
mjtNum kMaxDepth = 0.05;
|
||||
mjtNum pressure = depth / mju_max(kMaxDepth - depth, mjMINVAL);
|
||||
mju_scl3(force, normal, pressure);
|
||||
|
||||
// accumulate into forcesT (disjoint writes per taxel j)
|
||||
t->forcesT[0*ncon + j] += mju_dot3(force, normal);
|
||||
|
||||
if (has_frame) {
|
||||
mjtNum tang1[3] = {mesh_normal[normal_stride*j + 3],
|
||||
mesh_normal[normal_stride*j + 4],
|
||||
mesh_normal[normal_stride*j + 5]};
|
||||
mjtNum tang2[3] = {mesh_normal[normal_stride*j + 6],
|
||||
mesh_normal[normal_stride*j + 7],
|
||||
mesh_normal[normal_stride*j + 8]};
|
||||
mju_rotVecQuat(tang1, tang1, m->mesh_quat + 4 * mesh_id);
|
||||
mju_rotVecQuat(tang2, tang2, m->mesh_quat + 4 * mesh_id);
|
||||
t->forcesT[1*ncon + j] += mju_abs(mju_dot3(vel_rel, tang1));
|
||||
t->forcesT[2*ncon + j] += mju_abs(mju_dot3(vel_rel, tang2));
|
||||
}
|
||||
}
|
||||
}
|
||||
return NULL;
|
||||
}
|
||||
|
||||
|
||||
// apply cutoff to sensor i, clamping values in data buffer
|
||||
static void apply_cutoff(const mjModel* m, int i, mjtNum* data) {
|
||||
mjtNum cutoff = m->sensor_cutoff[i];
|
||||
@@ -1112,129 +1251,63 @@ static void mj_computeSensorAcc(const mjModel* m, mjData* d, int i, mjtNum* sens
|
||||
// all of the quadrature points are contact points
|
||||
int ncon = m->mesh_vertnum[mesh_id];
|
||||
|
||||
// get site frame
|
||||
mjtNum* geom_pos = d->geom_xpos + 3*geom_id;
|
||||
mjtNum* geom_mat = d->geom_xmat + 9*geom_id;
|
||||
|
||||
// allocate contact forces and positions
|
||||
mjtNum* forcesT = mj_stackAllocNum(d, ncon*3);
|
||||
mju_zero(forcesT, ncon*3);
|
||||
|
||||
// iterate over colliding geoms
|
||||
for (int g = 0; g < ncontact; g++) {
|
||||
int geom = contact_geom_ids[g];
|
||||
int body = m->geom_bodyid[geom];
|
||||
// threshold for parallelization (taxel count below which sequential is faster)
|
||||
const int kTactileParallelThreshold = 1000;
|
||||
|
||||
// get sdf plugin of the geoms
|
||||
int sdf_instance[2] = {-1, -1};
|
||||
mjtGeom geomtype[2] = {mjGEOM_SDF, mjGEOM_SPHERE};
|
||||
const mjpPlugin* sdf_ptr[2] = {NULL, NULL};
|
||||
if (m->geom_type[geom] == mjGEOM_SDF) {
|
||||
sdf_instance[0] = m->geom_plugin[geom];
|
||||
sdf_ptr[0] = mjc_getSDF(m, geom);
|
||||
} else if (m->geom_type[geom] == mjGEOM_MESH) {
|
||||
sdf_instance[0] = m->geom_dataid[geom];
|
||||
geomtype[0] = (mjtGeom)m->geom_type[geom];
|
||||
} else {
|
||||
sdf_instance[0] = geom;
|
||||
geomtype[0] = (mjtGeom)m->geom_type[geom];
|
||||
// parallel path: use threadpool to process taxel batches
|
||||
if (d->threadpool && ncon >= kTactileParallelThreshold) {
|
||||
int nthreads = mju_threadPoolNumberOfThreads((mjThreadPool*)d->threadpool);
|
||||
int batch_size = (ncon + nthreads - 1) / nthreads;
|
||||
int ntasks = (ncon + batch_size - 1) / batch_size;
|
||||
|
||||
mjTask* tasks = mjSTACKALLOC(d, ntasks, mjTask);
|
||||
mjTactileTaskArgs* task_args = mjSTACKALLOC(d, ntasks, mjTactileTaskArgs);
|
||||
|
||||
for (int t = 0; t < ntasks; t++) {
|
||||
task_args[t].m = m;
|
||||
task_args[t].d = d;
|
||||
task_args[t].sensor_id = i;
|
||||
task_args[t].mesh_id = mesh_id;
|
||||
task_args[t].geom_id = geom_id;
|
||||
task_args[t].parent_weld = parent_weld;
|
||||
task_args[t].ncontact = ncontact;
|
||||
task_args[t].nchannel = nchannel;
|
||||
task_args[t].contact_geom_ids = contact_geom_ids;
|
||||
task_args[t].start_taxel = t * batch_size;
|
||||
task_args[t].end_taxel = mju_min((t+1) * batch_size, ncon);
|
||||
task_args[t].forcesT = forcesT;
|
||||
|
||||
mju_defaultTask(&tasks[t]);
|
||||
tasks[t].func = tactile_taxel_batch;
|
||||
tasks[t].args = &task_args[t];
|
||||
mju_threadPoolEnqueue((mjThreadPool*)d->threadpool, &tasks[t]);
|
||||
}
|
||||
|
||||
// skip mesh geoms not having an octree
|
||||
if (geomtype[0] == mjGEOM_MESH &&
|
||||
m->mesh_octadr[m->geom_dataid[geom]] == -1) {
|
||||
continue;
|
||||
}
|
||||
|
||||
// set SDF parameters
|
||||
mjSDF geom_sdf;
|
||||
geom_sdf.id = &sdf_instance[0];
|
||||
geom_sdf.type = mjSDFTYPE_SINGLE;
|
||||
geom_sdf.plugin = &sdf_ptr[0];
|
||||
geom_sdf.geomtype = &geomtype[0];
|
||||
|
||||
// get forces in mesh coordinates
|
||||
int node = 0;
|
||||
float* mesh_vert = m->mesh_vert + 3*m->mesh_vertadr[mesh_id];
|
||||
float* mesh_normal = m->mesh_normal + 3*m->mesh_normaladr[mesh_id];
|
||||
|
||||
// check if mesh has frame or normal
|
||||
int has_frame = (m->mesh_normalnum[mesh_id] == 3 * m->mesh_vertnum[mesh_id]);
|
||||
int normal_stride = has_frame ? 9 : 3;
|
||||
|
||||
for (int j = 0; j < ncon; j++) {
|
||||
// position in site frame
|
||||
mjtNum pos[3] = {mesh_vert[3*j + 0], mesh_vert[3*j + 1], mesh_vert[3*j + 2]};
|
||||
|
||||
// position in global frame
|
||||
mjtNum xpos[3];
|
||||
mju_mulMatVec3(xpos, geom_mat, pos);
|
||||
mju_addTo3(xpos, geom_pos);
|
||||
|
||||
// position in other geom frame
|
||||
mjtNum lpos[3];
|
||||
mju_sub3(tmp, xpos, d->geom_xpos + 3*geom);
|
||||
mju_mulMatTVec3(lpos, d->geom_xmat + 9*geom, tmp);
|
||||
|
||||
// SDF plugins are in the original mesh frame
|
||||
if (sdf_ptr[0] != NULL) {
|
||||
mjtNum mesh_mat[9];
|
||||
mju_quat2Mat(mesh_mat, m->mesh_quat + 4 * m->geom_dataid[geom]);
|
||||
mju_mulMatVec3(lpos, mesh_mat, lpos);
|
||||
mju_addTo3(lpos, m->mesh_pos + 3 * m->geom_dataid[geom]);
|
||||
}
|
||||
|
||||
// compute distance
|
||||
mjtNum depth = mju_min(mjc_distance(m, d, &geom_sdf, lpos), 0);
|
||||
if (depth == 0) {
|
||||
node++;
|
||||
continue;
|
||||
}
|
||||
|
||||
// get velocity in global frame
|
||||
mjtNum vel_sensor[6], vel_other[6], vel_rel[3];
|
||||
mju_transformSpatial(
|
||||
vel_sensor, d->cvel + 6 * parent_weld, 0, xpos,
|
||||
d->subtree_com + 3 * m->body_rootid[parent_weld], NULL);
|
||||
mju_transformSpatial(
|
||||
vel_other, d->cvel + 6 * body, 0, d->geom_xpos + 3 * geom,
|
||||
d->subtree_com + 3 * m->body_rootid[body], NULL);
|
||||
mju_sub3(vel_rel, vel_sensor+3, vel_other+3);
|
||||
|
||||
// get normal
|
||||
mjtNum normal[3] = {mesh_normal[normal_stride*j + 0],
|
||||
mesh_normal[normal_stride*j + 1],
|
||||
mesh_normal[normal_stride*j + 2]};
|
||||
mju_rotVecQuat(normal, normal, m->mesh_quat + 4 * mesh_id);
|
||||
|
||||
// get contact force/torque
|
||||
mjtNum force[3];
|
||||
mjtNum kMaxDepth = 0.05;
|
||||
mjtNum pressure = depth / mju_max(kMaxDepth - depth, mjMINVAL);
|
||||
mju_scl3(force, normal, pressure);
|
||||
|
||||
// normal force
|
||||
forcesT[0*ncon + node] = mju_dot3(force, normal);
|
||||
|
||||
// tangent forces (require tangent frame)
|
||||
if (has_frame) {
|
||||
mjtNum tang1[3] = {mesh_normal[normal_stride*j + 3],
|
||||
mesh_normal[normal_stride*j + 4],
|
||||
mesh_normal[normal_stride*j + 5]};
|
||||
mjtNum tang2[3] = {mesh_normal[normal_stride*j + 6],
|
||||
mesh_normal[normal_stride*j + 7],
|
||||
mesh_normal[normal_stride*j + 8]};
|
||||
mju_rotVecQuat(tang1, tang1, m->mesh_quat + 4 * mesh_id);
|
||||
mju_rotVecQuat(tang2, tang2, m->mesh_quat + 4 * mesh_id);
|
||||
forcesT[1*ncon + node] = mju_abs(mju_dot3(vel_rel, tang1));
|
||||
forcesT[2*ncon + node] = mju_abs(mju_dot3(vel_rel, tang2));
|
||||
} else {
|
||||
forcesT[1*ncon + node] = 0;
|
||||
forcesT[2*ncon + node] = 0;
|
||||
}
|
||||
node++;
|
||||
for (int t = 0; t < ntasks; t++) {
|
||||
mju_taskJoin(&tasks[t]);
|
||||
}
|
||||
}
|
||||
// sequential path: call tactile_taxel_batch with full range
|
||||
else {
|
||||
mjTactileTaskArgs args;
|
||||
args.m = m;
|
||||
args.d = d;
|
||||
args.sensor_id = i;
|
||||
args.mesh_id = mesh_id;
|
||||
args.geom_id = geom_id;
|
||||
args.parent_weld = parent_weld;
|
||||
args.ncontact = ncontact;
|
||||
args.nchannel = nchannel;
|
||||
args.contact_geom_ids = contact_geom_ids;
|
||||
args.start_taxel = 0;
|
||||
args.end_taxel = ncon;
|
||||
args.forcesT = forcesT;
|
||||
tactile_taxel_batch(&args);
|
||||
}
|
||||
|
||||
// compute sensor output
|
||||
for (int c = 0; c < nchannel; c++) {
|
||||
|
||||
Reference in New Issue
Block a user