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_errmem.h"
|
||||||
#include "engine/engine_util_misc.h"
|
#include "engine/engine_util_misc.h"
|
||||||
#include "engine/engine_util_spatial.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);
|
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
|
// apply cutoff to sensor i, clamping values in data buffer
|
||||||
static void apply_cutoff(const mjModel* m, int i, mjtNum* data) {
|
static void apply_cutoff(const mjModel* m, int i, mjtNum* data) {
|
||||||
mjtNum cutoff = m->sensor_cutoff[i];
|
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
|
// all of the quadrature points are contact points
|
||||||
int ncon = m->mesh_vertnum[mesh_id];
|
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
|
// allocate contact forces and positions
|
||||||
mjtNum* forcesT = mj_stackAllocNum(d, ncon*3);
|
mjtNum* forcesT = mj_stackAllocNum(d, ncon*3);
|
||||||
mju_zero(forcesT, ncon*3);
|
mju_zero(forcesT, ncon*3);
|
||||||
|
|
||||||
// iterate over colliding geoms
|
// threshold for parallelization (taxel count below which sequential is faster)
|
||||||
for (int g = 0; g < ncontact; g++) {
|
const int kTactileParallelThreshold = 1000;
|
||||||
int geom = contact_geom_ids[g];
|
|
||||||
int body = m->geom_bodyid[geom];
|
|
||||||
|
|
||||||
// get sdf plugin of the geoms
|
// parallel path: use threadpool to process taxel batches
|
||||||
int sdf_instance[2] = {-1, -1};
|
if (d->threadpool && ncon >= kTactileParallelThreshold) {
|
||||||
mjtGeom geomtype[2] = {mjGEOM_SDF, mjGEOM_SPHERE};
|
int nthreads = mju_threadPoolNumberOfThreads((mjThreadPool*)d->threadpool);
|
||||||
const mjpPlugin* sdf_ptr[2] = {NULL, NULL};
|
int batch_size = (ncon + nthreads - 1) / nthreads;
|
||||||
if (m->geom_type[geom] == mjGEOM_SDF) {
|
int ntasks = (ncon + batch_size - 1) / batch_size;
|
||||||
sdf_instance[0] = m->geom_plugin[geom];
|
|
||||||
sdf_ptr[0] = mjc_getSDF(m, geom);
|
mjTask* tasks = mjSTACKALLOC(d, ntasks, mjTask);
|
||||||
} else if (m->geom_type[geom] == mjGEOM_MESH) {
|
mjTactileTaskArgs* task_args = mjSTACKALLOC(d, ntasks, mjTactileTaskArgs);
|
||||||
sdf_instance[0] = m->geom_dataid[geom];
|
|
||||||
geomtype[0] = (mjtGeom)m->geom_type[geom];
|
for (int t = 0; t < ntasks; t++) {
|
||||||
} else {
|
task_args[t].m = m;
|
||||||
sdf_instance[0] = geom;
|
task_args[t].d = d;
|
||||||
geomtype[0] = (mjtGeom)m->geom_type[geom];
|
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
|
for (int t = 0; t < ntasks; t++) {
|
||||||
if (geomtype[0] == mjGEOM_MESH &&
|
mju_taskJoin(&tasks[t]);
|
||||||
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++;
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
// 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
|
// compute sensor output
|
||||||
for (int c = 0; c < nchannel; c++) {
|
for (int c = 0; c < nchannel; c++) {
|
||||||
|
|||||||
Reference in New Issue
Block a user