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

Add octree to meshes of SDF geoms.

Browse Source 
PiperOrigin-RevId: 777634862
Change-Id: I3210c2f8d73fd30ceecb0a11abd1d82e086b3d5d
This commit is contained in:
Alessio Quaglino
2025-06-30 10:49:08 -07:00
committed by Copybara-Service
parent 106df98946
commit 5c1f1f9dba
17 changed files with 369 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/engine/engine_io.c
+4 -3
View File
@@ -460,7 +460,7 @@ static void freeModelBuffers(mjModel* m) {
// allocate and initialize mjModel structure
void mj_makeModel(mjModel** dest,
int nq, int nv, int nu, int na, int nbody, int nbvh,
int nbvhstatic, int nbvhdynamic, int njnt, int ngeom, int nsite, int ncam,
int nbvhstatic, int nbvhdynamic, int noct, int njnt, int ngeom, int nsite, int ncam,
int nlight, int nflex, int nflexnode, int nflexvert, int nflexedge, int nflexelem,
int nflexelemdata, int nflexelemedge, int nflexshelldata, int nflexevpair, int nflextexcoord,
int nmesh, int nmeshvert, int nmeshnormal, int nmeshtexcoord, int nmeshface,
@@ -499,6 +499,7 @@ void mj_makeModel(mjModel** dest,
m->nbvh = nbvh;
m->nbvhstatic = nbvhstatic;
m->nbvhdynamic = nbvhdynamic;
m->noct = noct;
m->njnt = njnt;
m->ngeom = ngeom;
m->nsite = nsite;
@@ -639,7 +640,7 @@ mjModel* mj_copyModel(mjModel* dest, const mjModel* src) {
if (!dest) {
mj_makeModel(&dest,
src->nq, src->nv, src->nu, src->na, src->nbody, src->nbvh,
src->nbvhstatic, src->nbvhdynamic, src->njnt, src->ngeom, src->nsite,
src->nbvhstatic, src->nbvhdynamic, src->noct, src->njnt, src->ngeom, src->nsite,
src->ncam, src->nlight, src->nflex, src->nflexnode, src->nflexvert, src->nflexedge,
src->nflexelem, src->nflexelemdata, src->nflexelemedge, src->nflexshelldata,
src->nflexevpair, src->nflextexcoord, src->nmesh, src->nmeshvert,
@@ -835,7 +836,7 @@ mjModel* mj_loadModelBuffer(const void* buffer, int buffer_sz) {
ints[42], ints[43], ints[44], ints[45], ints[46], ints[47], ints[48],
ints[49], ints[50], ints[51], ints[52], ints[53], ints[54], ints[55],
ints[56], ints[57], ints[58], ints[59], ints[60], ints[61], ints[62],
ints[63], ints[64], ints[65], ints[66], ints[67], ints[68]);
ints[63], ints[64], ints[65], ints[66], ints[67], ints[68], ints[69]);
if (!m || m->nbuffer != sizes[getnsize()-1]) {
mju_warning("Corrupted model, wrong size parameters");
mj_deleteModel(m);
src/engine/engine_io.h
+1 -1
View File
@@ -52,7 +52,7 @@ void mj_defaultStatistic(mjStatistic* stat);
// allocate mjModel
void mj_makeModel(mjModel** dest,
int nq, int nv, int nu, int na, int nbody, int nbvh, int nbvhstatic, int nbvhdynamic,
int nq, int nv, int nu, int na, int nbody, int nbvh, int nbvhstatic, int nbvhdynamic, int noct,
int njnt, int ngeom, int nsite, int ncam, int nlight, int nflex, int nflexnode, int nflexvert,
int nflexedge, int nflexelem, int nflexelemdata, int nflexelemedge, int nflexshelldata,
int nflexevpair, int nflextexcoord, int nmesh, int nmeshvert, int nmeshnormal,
src/engine/engine_vis_visualize.c
+38 -1
View File
@@ -763,7 +763,7 @@ void mjv_addGeoms(const mjModel* m, mjData* d, const mjvOption* vopt,
if (vopt->flags[mjVIS_MESHBVH]) {
for (int geomid = 0; geomid < m->ngeom; geomid++) {
int meshid = m->geom_dataid[geomid];
if (meshid == -1) {
if (m->geom_type[geomid] == mjGEOM_SDF || meshid == -1) {
continue;
}
@@ -806,6 +806,43 @@ void mjv_addGeoms(const mjModel* m, mjData* d, const mjvOption* vopt,
}
}
// mesh octree
category = mjCAT_DECOR;
objtype = mjOBJ_UNKNOWN;
if (vopt->flags[mjVIS_MESHBVH]) {
for (int geomid = 0; geomid < m->ngeom; geomid++) {
int meshid = m->geom_dataid[geomid];
if (m->geom_type[geomid] != mjGEOM_SDF || meshid == -1) {
continue;
}
for (int b = 0; b < m->mesh_octnum[meshid]; b++) {
int i = b + m->mesh_octadr[meshid];
START
if (m->oct_depth[i] != vopt->bvh_depth) {
continue;
}
// box color
const float* rgba = m->vis.rgba.bv;
// get xpos, xmat, size
const mjtNum* xpos = d->geom_xpos + 3 * geomid;
const mjtNum* xmat = d->geom_xmat + 9 * geomid;
const mjtNum *size = m->oct_aabb + 6*i + 3;
// offset xpos with aabb center (not always at geom origin)
const mjtNum *center = m->oct_aabb + 6*i;
mjtNum pos[3];
mju_mulMatVec3(pos, xmat, center);
mju_addTo3(pos, xpos);
mjv_initGeom(thisgeom, mjGEOM_LINEBOX, size, pos, xmat, rgba);
FINISH
}
}
}
// inertia
objtype = mjOBJ_BODY;
if (vopt->flags[mjVIS_INERTIA]) {
src/user/user_mesh.cc
+12
View File
@@ -139,6 +139,7 @@ mjCMesh::mjCMesh(mjCModel* _model, mjCDef* _def) {
maxhullvert_ = -1;
processed_ = false;
visual_ = true;
needoct_ = false;
// reset to default if given
if (_def) {
@@ -406,6 +407,7 @@ void mjCMesh::CacheMesh(mjCCache* cache, const mjResource* resource) {
}
mesh->tree_ = tree_;
mesh->face_aabb_ = face_aabb_;
mesh->octree_ = octree_;
// calculate estimated size of mesh
std::size_t size = sizeof(mjCMesh)
@@ -423,6 +425,7 @@ void mjCMesh::CacheMesh(mjCCache* cache, const mjResource* resource) {
+ (sizeof(double) * 18)
+ (sizeof(int) * ncenter)
+ tree_.Size()
+ octree_.Size()
+ (sizeof(double) * face_aabb_.size());
std::shared_ptr<const void> cached_data(mesh, +[] (const void* data) {
@@ -682,6 +685,14 @@ void mjCMesh::TryCompile(const mjVFS* vfs) {
}
}
// make octree
if (!needoct_) {
octree_.Clear();
} else if (octree_.Nodes().empty()) {
octree_.SetFace(vert_, face_);
octree_.CreateOctree(aamm_);
}
// close resource
if (resource_ != nullptr) {
mju_closeResource(resource_);
@@ -1073,6 +1084,7 @@ bool mjCMesh::LoadCachedMesh(mjCCache *cache, const mjResource* resource) {
}
tree_ = mesh->tree_;
face_aabb_ = mesh->face_aabb_;
octree_ = mesh->octree_;
return true;
};
src/user/user_model.cc
+17 -2
View File
@@ -963,6 +963,7 @@ void mjCModel::Clear() {
nbvh = 0;
nbvhstatic = 0;
nbvhdynamic = 0;
noct = 0;
njnt = 0;
ngeom = 0;
nsite = 0;
@@ -1726,6 +1727,7 @@ void mjCModel::IndexAssets(bool discard) {
((mjCMesh*)mesh)->SetNotVisual(); // reset to true by mesh->Compile()
}
geom->mesh = (discard && geom->visual_) ? nullptr : (mjCMesh*)mesh;
static_cast<mjCMesh*>(mesh)->needoct_ |= geom->spec.type == mjGEOM_SDF;
} else {
throw mjCError(geom, "mesh '%s' not found in geom %d", geom->get_meshname().c_str(), i);
}
@@ -1936,6 +1938,7 @@ void mjCModel::SetSizes() {
}
for (int i=0; i < nmesh; i++) {
nbvhstatic += meshes_[i]->tree().Nbvh();
noct += meshes_[i]->octree().NumNodes();
}
for (int i=0; i < nflex; i++) {
nbvhdynamic += flexes_[i]->tree.Nbvh();
@@ -2943,7 +2946,7 @@ int mjCModel::CountNJmom(const mjModel* m) {
// copy objects outside kinematic tree
void mjCModel::CopyObjects(mjModel* m) {
int adr, bone_adr, vert_adr, node_adr, normal_adr, face_adr, texcoord_adr;
int adr, bone_adr, vert_adr, node_adr, normal_adr, face_adr, texcoord_adr, oct_adr;
int edge_adr, elem_adr, elemdata_adr, elemedge_adr, shelldata_adr, evpair_adr;
int bonevert_adr, graph_adr, data_adr, bvh_adr;
int poly_adr, polymap_adr, polyvert_adr;
@@ -2963,6 +2966,7 @@ void mjCModel::CopyObjects(mjModel* m) {
}
// meshes
oct_adr = 0;
vert_adr = 0;
normal_adr = 0;
texcoord_adr = 0;
@@ -2989,6 +2993,8 @@ void mjCModel::CopyObjects(mjModel* m) {
m->mesh_graphadr[i] = (pme->szgraph() ? graph_adr : -1);
m->mesh_bvhnum[i] = pme->tree().Nbvh();
m->mesh_bvhadr[i] = pme->tree().Nbvh() ? bvh_adr : -1;
m->mesh_octnum[i] = pme->octree().NumNodes();
m->mesh_octadr[i] = pme->octree().NumNodes() ? oct_adr : -1;
mjuu_copyvec(&m->mesh_scale[3 * i], pme->Scale(), 3);
mjuu_copyvec(&m->mesh_pos[3 * i], pme->GetPosPtr(), 3);
mjuu_copyvec(&m->mesh_quat[4 * i], pme->GetQuatPtr(), 4);
@@ -3023,6 +3029,14 @@ void mjCModel::CopyObjects(mjModel* m) {
}
}
// copy octree data
if (pme->octree().NumNodes()) {
int n_oct = pme->octree().NumNodes();
memcpy(m->oct_aabb + 6*oct_adr, pme->octree().Nodes().data(), 6*n_oct*sizeof(mjtNum));
memcpy(m->oct_child + 8*oct_adr, pme->octree().Child().data(), 8*n_oct*sizeof(int));
memcpy(m->oct_depth + oct_adr, pme->octree().Level().data(), n_oct*sizeof(int));
}
// advance counters
poly_adr += pme->npolygon();
polyvert_adr += pme->npolygonvert();
@@ -3033,6 +3047,7 @@ void mjCModel::CopyObjects(mjModel* m) {
face_adr += pme->nface();
graph_adr += pme->szgraph();
bvh_adr += pme->tree().Nbvh();
oct_adr += pme->octree().NumNodes();
}
// flexes
@@ -4523,7 +4538,7 @@ void mjCModel::TryCompile(mjModel*& m, mjData*& d, const mjVFS* vfs) {
// create low-level model
mj_makeModel(&m,
nq, nv, nu, na, nbody, nbvh, nbvhstatic, nbvhdynamic, njnt, ngeom, nsite,
nq, nv, nu, na, nbody, nbvh, nbvhstatic, nbvhdynamic, noct, njnt, ngeom, nsite,
ncam, nlight, nflex, nflexnode, nflexvert, nflexedge, nflexelem,
nflexelemdata, nflexelemedge, nflexshelldata, nflexevpair, nflextexcoord,
nmesh, nmeshvert, nmeshnormal, nmeshtexcoord, nmeshface, nmeshgraph, nmeshpoly,
src/user/user_model.h
+1
View File
@@ -87,6 +87,7 @@ class mjCModel_ : public mjsElement {
int nbvh; // number of total boundary volume hierarchies
int nbvhstatic; // number of static boundary volume hierarchies
int nbvhdynamic; // number of dynamic boundary volume hierarchies
int noct; // number of total octree cells
int nflexnode; // number of nodes in all flexes
int nflexvert; // number of vertices in all flexes
int nflexedge; // number of edges in all flexes
src/user/user_objects.cc
+109 -4
View File
@@ -404,9 +404,15 @@ mjCBoundingVolumeHierarchy::AddBoundingVolume(const int* id, int contype, int co
// create bounding volume hierarchy
void mjCBoundingVolumeHierarchy::CreateBVH() {
std::vector<BVElement> elements;
Make(elements);
MakeBVH(elements.begin(), elements.end());
}
void mjCBoundingVolumeHierarchy::Make(std::vector<BVElement>& elements) {
// precompute the positions of each element in the hierarchy's axes, and drop
// visual-only elements.
std::vector<BVElement> elements;
elements.reserve(bvleaf_.size());
double qinv[4] = {iquat_[0], -iquat_[1], -iquat_[2], -iquat_[3]};
for (int i = 0; i < bvleaf_.size(); i++) {
@@ -420,9 +426,9 @@ void mjCBoundingVolumeHierarchy::CreateBVH() {
elements.push_back(std::move(element));
}
}
MakeBVH(elements.begin(), elements.end());
}
// compute bounding volume hierarchy
int mjCBoundingVolumeHierarchy::MakeBVH(
std::vector<BVElement>::iterator elements_begin,
@@ -546,10 +552,109 @@ int mjCBoundingVolumeHierarchy::MakeBVH(
return index;
}
//------------------------- class mjCOctree implementation --------------------------------------------
void mjCOctree::SetFace(const std::vector<double>& vert, const std::vector<int>& face) {
for (int i = 0; i < face.size(); i += 3) {
std::array<double, 3> v0 = {vert[3*face[i+0]], vert[3*face[i+0]+1], vert[3*face[i+0]+2]};
std::array<double, 3> v1 = {vert[3*face[i+1]], vert[3*face[i+1]+1], vert[3*face[i+1]+2]};
std::array<double, 3> v2 = {vert[3*face[i+2]], vert[3*face[i+2]+1], vert[3*face[i+2]+2]};
face_.push_back({v0, v1, v2});
}
}
// TODO: use the same code as mjCBoundingVolumeHierarchy::Make()
void mjCOctree::Make(std::vector<Triangle>& elements) {
// rotate triangles to the body inertial frame
elements.assign(face_.size(), {{{0}}});
double qinv[4] = {iquat_[0], -iquat_[1], -iquat_[2], -iquat_[3]};
for (int i = 0; i < face_.size(); i++) {
for (int j = 0; j < 3; j++) {
double vert[3] = {face_[i][j][0] - ipos_[0],
face_[i][j][1] - ipos_[1],
face_[i][j][2] - ipos_[2]};
mjuu_rotVecQuat(elements[i][j].data(), vert, qinv);
}
}
}
void mjCOctree::CreateOctree(const double aamm[6]) {
std::vector<Triangle> elements;
Make(elements);
std::vector<Triangle*> elements_ptrs(elements.size());
std::transform(elements.begin(), elements.end(), elements_ptrs.begin(),
[](Triangle& triangle) { return &triangle; });
MakeOctree(elements_ptrs, aamm);
}
static bool boxTriangle(const Triangle& element, const double aamm[6]) {
for (int i = 0; i < 3; i++) {
if (element[0][i] < aamm[i] && element[1][i] < aamm[i] && element[2][i] < aamm[i]) {
return false;
}
int j = i + 3;
if (element[0][i] > aamm[j] && element[1][i] > aamm[j] && element[2][i] > aamm[j]) {
return false;
}
}
// TODO: add additionally separating axis tests
return true;
}
int mjCOctree::MakeOctree(const std::vector<Triangle*>& elements, const double aamm[6], int lev) {
level_.push_back(lev);
// create a new node
int index = nnode_++;
double aabb[6] = {(aamm[0] + aamm[3]) / 2, (aamm[1] + aamm[4]) / 2, (aamm[2] + aamm[5]) / 2,
(aamm[3] - aamm[0]) / 2, (aamm[4] - aamm[1]) / 2, (aamm[5] - aamm[2]) / 2};
for (int i = 0; i < 6; i++) {
node_.push_back(aabb[i]);
}
for (int i = 0; i < 8; i++) {
child_.push_back(-1);
}
// find all triangles that intersect the current box
std::vector<Triangle*> colliding;
for (auto* element : elements) {
if (boxTriangle(*element, aamm)) {
colliding.push_back(element);
}
}
// return if the box is empty
if (colliding.empty() || lev >= 6) {
return index;
}
// split the box into 8 sub-boxes
double new_aamm[8][6];
for (int i = 0; i < 8; i++) {
new_aamm[i][0] = aabb[0] + aabb[3] * (i & 1 ? -1 : 0);
new_aamm[i][1] = aabb[1] + aabb[4] * (i & 2 ? -1 : 0);
new_aamm[i][2] = aabb[2] + aabb[5] * (i & 4 ? -1 : 0);
new_aamm[i][3] = aabb[0] + aabb[3] * (i & 1 ? 0 : 1);
new_aamm[i][4] = aabb[1] + aabb[4] * (i & 2 ? 0 : 1);
new_aamm[i][5] = aabb[2] + aabb[5] * (i & 4 ? 0 : 1);
}
// recursive calls to create sub-boxes
for (int i = 0; i < 8; i++) {
child_[8*index + i] = MakeOctree(colliding, new_aamm[i], lev + 1);
}
return index;
}
//------------------------- class mjCDef implementation --------------------------------------------
// constructor
mjCDef::mjCDef() {
name.clear();
src/user/user_objects.h
+50 -1
View File
@@ -206,13 +206,54 @@ class mjCBoundingVolumeHierarchy : public mjCBoundingVolumeHierarchy_ {
// position of the element in the BVH axes
double lpos[3];
};
void Make(std::vector<BVElement>& elements);
int MakeBVH(std::vector<BVElement>::iterator elements_begin,
std::vector<BVElement>::iterator elements_end, int lev = 0);
};
//------------------------- class mjCOctree --------------------------------------------------------
typedef std::array<std::array<double, 3>, 3> Triangle;
struct mjCOctree_ {
int nnode_ = 0;
std::vector<int> child_; // children of each node (nnode x 8)
std::vector<double> node_; // bounding boxes (nnode x 6)
std::vector<int> level_; // levels of each node (nnode x 1)
std::vector<Triangle> face_; // mesh faces (nface x 3)
double ipos_[3] = {0, 0, 0};
double iquat_[4] = {1, 0, 0, 0};
};
class mjCOctree : public mjCOctree_ {
public:
void CreateOctree(const double aamm[6]);
int NumNodes() const { return nnode_; }
const std::vector<int>& Child() const { return child_; }
const std::vector<double>& Nodes() const { return node_; }
const std::vector<int>& Level() const { return level_; }
void SetFace(const std::vector<double>& vert, const std::vector<int>& face);
int Size() const {
return sizeof(int) * child_.size() + sizeof(double) * node_.size() +
sizeof(int) * level_.size() + sizeof(Triangle) * face_.size();
}
void Clear() {
child_.clear();
node_.clear();
level_.clear();
face_.clear();
}
private:
void Make(std::vector<Triangle>& elements);
int MakeOctree(const std::vector<Triangle*>& elements, const double aamm[6], int lev = 0);
};
//------------------------- class mjCBase ----------------------------------------------------------
// Generic functionality for all derived classes
@@ -965,6 +1006,7 @@ class mjCMesh_ : public mjCBase {
std::vector<int> spec_facetexcoord_;
// used by the compiler
bool needoct_; // needs octree
bool visual_; // true: the mesh is only visual
std::vector< std::pair<int, int> > halfedge_; // half-edge data
@@ -985,9 +1027,13 @@ class mjCMesh_ : public mjCBase {
bool needhull_; // needs convex hull for collisions
int maxhullvert_; // max vertex count of convex hull
// bounding volume hierarchy tree
mjCBoundingVolumeHierarchy tree_; // bounding volume hierarchy
std::vector<double> face_aabb_; // bounding boxes of all faces
// octree
mjCOctree octree_; // octree of the mesh
// paths stored during model attachment
mujoco::user::FilePath modelfiledir_;
mujoco::user::FilePath meshdir_;
@@ -1061,6 +1107,9 @@ class mjCMesh: public mjCMesh_, private mjsMesh {
// bounding volume hierarchy tree
const mjCBoundingVolumeHierarchy& tree() { return tree_; }
// octree
const mjCOctree& octree() { return octree_; }
void Compile(const mjVFS* vfs); // compiler
double* GetPosPtr(); // get position
double* GetQuatPtr(); // get orientation