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

Move tree construction to helper class.

Browse Source 
PiperOrigin-RevId: 525171924
Change-Id: I28b572fa09a2ccc029e56f9a84808ac53a4e6ac9
This commit is contained in:
Alessio Quaglino
2023-04-18 09:41:55 -07:00
committed by Copybara-Service
parent 94c4f72907
commit b25728cc2e
3 changed files with 199 additions and 170 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/user/user_model.cc
+8 -8
View File
@@ -928,7 +928,7 @@ void mjCModel::SetSizes(void) {
// nbvh
for (int i=0; i<nbody; i++) {
nbvh += bodies[i]->nbvh;
nbvh += bodies[i]->tree.nbvh;
}
// nmeshvert, nmeshface, nmeshtexcoord, nmeshgraph
@@ -1346,14 +1346,14 @@ void mjCModel::CopyTree(mjModel* m) {
// bounding volume hierarchy
m->body_bvhadr[i] = (!pb->geoms.empty() ? bvh_adr : -1);
m->body_bvhnum[i] = pb->nbvh;
if (pb->nbvh) {
memcpy(m->bvh_aabb + 6*bvh_adr, pb->bvh.data(), 6*pb->nbvh*sizeof(mjtNum));
memcpy(m->bvh_child + 2*bvh_adr, pb->child.data(), 2*pb->nbvh*sizeof(int));
memcpy(m->bvh_geomid + bvh_adr, pb->nodeid.data(), pb->nbvh*sizeof(int));
memcpy(m->bvh_depth + bvh_adr, pb->level.data(), pb->nbvh*sizeof(int));
m->body_bvhnum[i] = pb->tree.nbvh;
if (pb->tree.nbvh) {
memcpy(m->bvh_aabb + 6*bvh_adr, pb->tree.bvh.data(), 6*pb->tree.nbvh*sizeof(mjtNum));
memcpy(m->bvh_child + 2*bvh_adr, pb->tree.child.data(), 2*pb->tree.nbvh*sizeof(int));
memcpy(m->bvh_geomid + bvh_adr, pb->tree.nodeid.data(), pb->tree.nbvh*sizeof(int));
memcpy(m->bvh_depth + bvh_adr, pb->tree.level.data(), pb->tree.nbvh*sizeof(int));
}
bvh_adr += pb->nbvh;
bvh_adr += pb->tree.nbvh;
// count free joints
int cntfree = 0;
src/user/user_objects.cc
+168 -154
View File
@@ -263,6 +263,172 @@ const char* mjCAlternative::Set(double* quat, double* inertia,
//------------------------- class mjCTree implementation -------------------------------------------
// constructor
mjCBoundingVolumeHierarchy::mjCBoundingVolumeHierarchy() {
nbvh = 0;
mjuu_setvec(ipos_, 0, 0, 0);
mjuu_setvec(iquat_, 1, 0, 0, 0);
}
// assign position and orientation
void mjCBoundingVolumeHierarchy::Set(mjtNum ipos_element[3], mjtNum iquat_element[4]) {
mjuu_copyvec(ipos_, ipos_element, 3);
mjuu_copyvec(iquat_, iquat_element, 4);
}
// compute bounding volume hierarchy
int mjCBoundingVolumeHierarchy::MakeBVH(std::vector<mjCGeom *>& elements, int lev) {
int nelements = elements.size();
mjtNum AABB[6] = {mjMAXVAL, mjMAXVAL, mjMAXVAL, -mjMAXVAL, -mjMAXVAL, -mjMAXVAL};
// inverse transformation
mjtNum qinv[4] = {iquat_[0], -iquat_[1], -iquat_[2], -iquat_[3]};
for (int i=0; i<nelements; i++) {
// skip visual objects
if (elements[i]->conaffinity==0 && elements[i]->contype==0) {
continue;
}
// transform aabb representation
mjtNum aabb[6] = {elements[i]->aabb[0] - elements[i]->aabb[3],
elements[i]->aabb[1] - elements[i]->aabb[4],
elements[i]->aabb[2] - elements[i]->aabb[5],
elements[i]->aabb[0] + elements[i]->aabb[3],
elements[i]->aabb[1] + elements[i]->aabb[4],
elements[i]->aabb[2] + elements[i]->aabb[5]};
// update node AABB
for (int v=0; v<8; v++) {
mjtNum vert[3], box[3];
vert[0] = (v&1 ? aabb[3] : aabb[0]);
vert[1] = (v&2 ? aabb[4] : aabb[1]);
vert[2] = (v&4 ? aabb[5] : aabb[2]);
// rotate to the body inertial frame
mju_rotVecQuat(box, vert, elements[i]->quat);
box[0] += elements[i]->pos[0] - ipos_[0];
box[1] += elements[i]->pos[1] - ipos_[1];
box[2] += elements[i]->pos[2] - ipos_[2];
mju_rotVecQuat(vert, box, qinv);
AABB[0] = mjMIN(AABB[0], vert[0]);
AABB[1] = mjMIN(AABB[1], vert[1]);
AABB[2] = mjMIN(AABB[2], vert[2]);
AABB[3] = mjMAX(AABB[3], vert[0]);
AABB[4] = mjMAX(AABB[4], vert[1]);
AABB[5] = mjMAX(AABB[5], vert[2]);
}
}
// store current index
int index = nbvh++;
child.push_back(-1);
child.push_back(-1);
nodeid.push_back(-1);
level.push_back(lev);
// transform representation
mjtNum center[] = {(AABB[3] + AABB[0]) / 2, (AABB[4] + AABB[1]) / 2,
(AABB[5] + AABB[2]) / 2};
mjtNum size[] = {(AABB[3] - AABB[0]) / 2, (AABB[4] - AABB[1]) / 2,
(AABB[5] - AABB[2]) / 2};
// store bounding box of the current node
for (int i=0; i<3; i++) {
bvh.push_back(center[i]);
}
for (int i=0; i<3; i++) {
bvh.push_back(size[i]);
}
// leaf node, return
if (nelements==1) {
for (int i=0; i<2; i++) {
child[2*index+i] = -1;
}
nodeid[index] = elements[0]->id;
return index;
}
// find longest axis for splitting the bounding box
mjtNum edges[3] = { AABB[3]-AABB[0], AABB[4]-AABB[1], AABB[5]-AABB[2] };
int axis = edges[0] > edges[1] ? 0 : 1;
axis = edges[axis] > edges[2] ? axis : 2;
// find median along the axis
std::vector<mjtNum> pos(nelements);
for (int i=0; i<nelements; i++) {
// get position in the body inertial frame
mjtNum vert[3] = {elements[i]->pos[0] - ipos_[0],
elements[i]->pos[1] - ipos_[1],
elements[i]->pos[2] - ipos_[2]};
mjtNum lpos[3];
mju_rotVecQuat(lpos, vert, qinv);
pos[i] = lpos[axis];
}
auto m = pos.size()/2;
std::nth_element(pos.begin(), pos.begin() + m, pos.end());
mjtNum threshold = pos[m];
// split using median
std::vector<mjCGeom *> left;
std::vector<mjCGeom *> right;
int skipped = 0;
for (int i=0; i<nelements; i++) {
// get position in the body inertial frame
mjtNum vert[3] = {elements[i]->pos[0] - ipos_[0],
elements[i]->pos[1] - ipos_[1],
elements[i]->pos[2] - ipos_[2]};
mjtNum lpos[3];
mju_rotVecQuat(lpos, vert, qinv);
// skip visual objects
if (elements[i]->conaffinity==0 && elements[i]->contype==0) {
skipped++;
continue;
}
if (lpos[axis] < threshold) {
left.push_back(elements[i]);
} else if (lpos[axis] > threshold) {
right.push_back(elements[i]);
} else {
if (left.size() < right.size()) left.push_back(elements[i]);
else right.push_back(elements[i]);
}
}
// recursive calls
if (!left.empty()) {
child[2*index+0] = MakeBVH(left, lev+1);
}
if (!right.empty()) {
child[2*index+1] = MakeBVH(right, lev+1);
}
// SHOULD NOT OCCUR
if (left.size()+right.size()+skipped != nelements) {
mju_error("some elements were lost, body=%s parent=%d children=%lu",
name_.c_str(), nelements, left.size()+right.size()+skipped);
}
if (child[2*index+0]==-1 && child[2*index+1]==-1 && !skipped) {
mju_error("this should have been a leaf, body=%s nelements=%d",
name_.c_str(), nelements);
}
return index;
}
//------------------------- class mjCDef implementation --------------------------------------------
// constructor
@@ -331,7 +497,6 @@ mjCBody::mjCBody(mjCModel* _model) {
subtreedofs = 0;
gravcomp = 0;
userdata.clear();
nbvh = 0;
// plugin variables
is_plugin = false;
@@ -346,10 +511,6 @@ mjCBody::mjCBody(mjCModel* _model) {
sites.clear();
cameras.clear();
lights.clear();
bvh.clear();
child.clear();
nodeid.clear();
level.clear();
}
@@ -659,154 +820,6 @@ void mjCBody::MakeInertialExplicit() {
}
// compute bounding volume hierarchy
int mjCBody::MakeBVH(std::vector<mjCGeom *>& elements, int lev) {
int nelements = elements.size();
mjtNum AABB[6] = {mjMAXVAL, mjMAXVAL, mjMAXVAL, -mjMAXVAL, -mjMAXVAL, -mjMAXVAL};
// inverse transformation
mjtNum qinv[4] = {iquat[0], -iquat[1], -iquat[2], -iquat[3]};
for (int i=0; i<nelements; i++) {
// skip visual objects
if (elements[i]->conaffinity==0 && elements[i]->contype==0) {
continue;
}
// transform aabb representation
mjtNum aabb[6] = {elements[i]->aabb[0] - elements[i]->aabb[3],
elements[i]->aabb[1] - elements[i]->aabb[4],
elements[i]->aabb[2] - elements[i]->aabb[5],
elements[i]->aabb[0] + elements[i]->aabb[3],
elements[i]->aabb[1] + elements[i]->aabb[4],
elements[i]->aabb[2] + elements[i]->aabb[5]};
// update node AABB
for (int v=0; v<8; v++) {
mjtNum vert[3], box[3];
vert[0] = (v&1 ? aabb[3] : aabb[0]);
vert[1] = (v&2 ? aabb[4] : aabb[1]);
vert[2] = (v&4 ? aabb[5] : aabb[2]);
// rotate to the body inertial frame
mju_rotVecQuat(box, vert, elements[i]->quat);
box[0] += elements[i]->pos[0] - ipos[0];
box[1] += elements[i]->pos[1] - ipos[1];
box[2] += elements[i]->pos[2] - ipos[2];
mju_rotVecQuat(vert, box, qinv);
AABB[0] = mjMIN(AABB[0], vert[0]);
AABB[1] = mjMIN(AABB[1], vert[1]);
AABB[2] = mjMIN(AABB[2], vert[2]);
AABB[3] = mjMAX(AABB[3], vert[0]);
AABB[4] = mjMAX(AABB[4], vert[1]);
AABB[5] = mjMAX(AABB[5], vert[2]);
}
}
// store current index
int index = nbvh++;
child.push_back(-1);
child.push_back(-1);
nodeid.push_back(-1);
level.push_back(lev);
// transform representation
mjtNum center[] = {(AABB[3] + AABB[0]) / 2, (AABB[4] + AABB[1]) / 2,
(AABB[5] + AABB[2]) / 2};
mjtNum size[] = {(AABB[3] - AABB[0]) / 2, (AABB[4] - AABB[1]) / 2,
(AABB[5] - AABB[2]) / 2};
// store bounding box of the current node
for (int i=0; i<3; i++) {
bvh.push_back(center[i]);
}
for (int i=0; i<3; i++) {
bvh.push_back(size[i]);
}
// leaf node, return
if (nelements==1) {
for (int i=0; i<2; i++) {
child[2*index+i] = -1;
}
nodeid[index] = elements[0]->id;
return index;
}
// find longest axis for splitting the bounding box
mjtNum edges[3] = { AABB[3]-AABB[0], AABB[4]-AABB[1], AABB[5]-AABB[2] };
int axis = edges[0] > edges[1] ? 0 : 1;
axis = edges[axis] > edges[2] ? axis : 2;
// find median along the axis
std::vector<mjtNum> pos(nelements);
for (int i=0; i<nelements; i++) {
// get position in the body inertial frame
mjtNum vert[3] = {elements[i]->pos[0] - ipos[0],
elements[i]->pos[1] - ipos[1],
elements[i]->pos[2] - ipos[2]};
mjtNum lpos[3];
mju_rotVecQuat(lpos, vert, qinv);
pos[i] = lpos[axis];
}
auto m = pos.size()/2;
std::nth_element(pos.begin(), pos.begin() + m, pos.end());
mjtNum threshold = pos[m];
// split using median
std::vector<mjCGeom *> left;
std::vector<mjCGeom *> right;
int skipped = 0;
for (int i=0; i<nelements; i++) {
// get position in the body inertial frame
mjtNum vert[3] = {elements[i]->pos[0] - ipos[0],
elements[i]->pos[1] - ipos[1],
elements[i]->pos[2] - ipos[2]};
mjtNum lpos[3];
mju_rotVecQuat(lpos, vert, qinv);
// skip visual objects
if (elements[i]->conaffinity==0 && elements[i]->contype==0) {
skipped++;
continue;
}
if (lpos[axis] < threshold) {
left.push_back(elements[i]);
} else if (lpos[axis] > threshold) {
right.push_back(elements[i]);
} else {
if (left.size() < right.size()) left.push_back(elements[i]);
else right.push_back(elements[i]);
}
}
// recursive calls
if (!left.empty()) {
child[2*index+0] = MakeBVH(left, lev+1);
}
if (!right.empty()) {
child[2*index+1] = MakeBVH(right, lev+1);
}
// SHOULD NOT OCCUR
if (left.size()+right.size()+skipped != nelements) {
throw mjCError(this, "some elements were lost, body=%s parent=%d children=%d",
name.c_str(), nelements, left.size()+right.size()+skipped);
}
if (child[2*index+0]==-1 && child[2*index+1]==-1 && !skipped) {
throw mjCError(this, "this should have been a leaf, body=%s nelements=%d",
name.c_str(), nelements);
}
return index;
}
// compiler
void mjCBody::Compile(void) {
// resize userdata
@@ -917,7 +930,8 @@ void mjCBody::Compile(void) {
// compute bounding volume hierarchy
if (!geoms.empty()) {
MakeBVH(geoms, 0);
tree.Set(ipos, iquat);
tree.MakeBVH(geoms);
}
// compile all joints, count dofs
src/user/user_objects.h
+23 -8
View File
@@ -114,6 +114,27 @@ class mjCAlternative {
};
// bounding volume hierarchy
class mjCBoundingVolumeHierarchy {
public:
mjCBoundingVolumeHierarchy();
int nbvh;
std::vector<mjtNum> bvh; // bounding boxes (nbvh x 6)
std::vector<int> child; // children of each node (nbvh x 2)
std::vector<int> nodeid; // id of the geom contained by the node (nbvh x 1)
std::vector<int> level; // levels of each node (nbvh x 1)
int MakeBVH(std::vector<mjCGeom *>&, int lev = 0); // make bounding volume hierarchy
void Set(mjtNum ipos_element[3], mjtNum iquat_element[4]);
private:
std::string name_;
double ipos_[3];
double iquat_[4];
};
//------------------------- class mjCBase ----------------------------------------------------------
// Generic functionality for all derived classes
@@ -213,14 +234,7 @@ class mjCBody : public mjCBase {
int lastdof; // id of last dof
int subtreedofs; // number of dofs in subtree, including self
int MakeBVH(std::vector<mjCGeom *>&, int lev); // make bounding volume hierarchy
int nbvh;
std::vector<mjtNum> bvh; // bounding volume hierarchy
std::vector<int> child; // children of bvh nodes
std::vector<int> nodeid; // id of the geom contained by the node
std::vector<int> level; // levels of bvh
mjCBoundingVolumeHierarchy tree; // bounding volume hierarchy
// objects allocated by Add functions
std::vector<mjCBody*> bodies; // child bodies
@@ -292,6 +306,7 @@ class mjCGeom : public mjCBase {
friend class mjCModel;
friend class mjXWriter;
friend class mjXURDF;
friend class mjCBoundingVolumeHierarchy;
public:
double GetVolume(void); // compute geom volume