From 56e98cc16ccc1410e788054ea90aada390b6e27c Mon Sep 17 00:00:00 2001
From: Kyle Bayes <kylebayes@google.com>
Date: Thu, 16 Apr 2026 03:24:28 -0700
Subject: [PATCH] Improve AABB bounds in makeAAMM.

MODEL           BEFORE(ns)   AFTER(ns)
boxmesh.xml     874645       744822
box.xml         808226       786226
ellipsoid.xml   1477609      1436058
mixed.xml       1492063      1404848

PiperOrigin-RevId: 900626282
Change-Id: I3df5a6cd8a4cb2e3b9ded1b40a5310a4c14ec624
---
 src/engine/engine_collision_driver.c | 32 ++++++++++++++++++++++------
 src/engine/engine_inline.h           | 15 +++++++++++++
 src/engine/engine_util_misc.c        | 12 ++---------
 3 files changed, 43 insertions(+), 16 deletions(-)

diff --git a/src/engine/engine_collision_driver.c b/src/engine/engine_collision_driver.c
index eb092942..b77164f6 100644
--- a/src/engine/engine_collision_driver.c
+++ b/src/engine/engine_collision_driver.c
@@ -27,6 +27,7 @@
 #include "engine/engine_collision_sdf.h"
 #include "engine/engine_core_constraint.h"
 #include "engine/engine_core_util.h"
+#include "engine/engine_inline.h"
 #include "engine/engine_macro.h"
 #include "engine/engine_memory.h"
 #include "engine/engine_sort.h"
@@ -949,6 +950,7 @@ static void makeAAMM(const mjModel* m, mjData* d,
                      mjtNum* x_max, mjtNum* y_max, mjtNum* z_max,
                      int bf, const mjtNum* frame) {
   mjtNum aamm[6];
+  mjtNum override_margin = mjENABLED(mjENBL_OVERRIDE) ? 0.5 * m->opt.o_margin : 0;
 
   // body
   if (bf < m->nbody) {
@@ -958,14 +960,32 @@ static void makeAAMM(const mjModel* m, mjData* d,
     // process all body geoms (body is collidable, should have geoms)
     for (int i=0; i < body_geomnum; i++) {
       int geom = m->body_geomadr[body]+i;
-      mjtNum margin = mjENABLED(mjENBL_OVERRIDE) ? 0.5*m->opt.o_margin : m->geom_margin[geom];
+      mjtNum margin = override_margin ? override_margin : m->geom_margin[geom];
       mjtNum _aamm[6];
 
-      // set _aamm for this geom
+      const mjtNum* aabb = m->geom_aabb + 6*geom;
+      const mjtNum* size = m->geom_aabb + 6*geom + 3;
+      const mjtNum* xpos = d->geom_xpos + 3*geom;
+      const mjtNum* xmat = d->geom_xmat + 9*geom;
+
+      // compute center in global coordinates
+      mjtNum pos[3];
+      mji_mulMatVec3(pos, xmat, aabb);
+      mju_addTo3(pos, xpos);
+
+      mjtNum axis[9];
+      mji_transpose3(axis, xmat);
+      mjtNum r_half = m->geom_rbound[geom];
+
       for (int j=0; j < 3; j++) {
-        mjtNum cen = mju_dot3(d->geom_xpos+3*geom, frame+3*j);
-        _aamm[j]   = cen - m->geom_rbound[geom] - margin;
-        _aamm[j+3] = cen + m->geom_rbound[geom] + margin;
+        const mjtNum* frame_j = frame + 3*j;
+        mjtNum aabb_cen = mju_dot3(pos, frame_j);
+        mjtNum aabb_half = mju_abs(size[0] * mju_dot3(axis + 0, frame_j))
+                         + mju_abs(size[1] * mju_dot3(axis + 3, frame_j))
+                         + mju_abs(size[2] * mju_dot3(axis + 6, frame_j));
+        mjtNum r_cen = mju_dot3(xpos, frame_j);
+        _aamm[j + 0] = mju_max(r_cen - r_half, aabb_cen - aabb_half) - margin;
+        _aamm[j + 3] = mju_min(r_cen + r_half, aabb_cen + aabb_half) + margin;
       }
 
       // update body aamm
@@ -1006,7 +1026,7 @@ static void makeAAMM(const mjModel* m, mjData* d,
     }
 
     // correct for flex radius and margin
-    mjtNum margin = mjENABLED(mjENBL_OVERRIDE) ? 0.5*m->opt.o_margin : m->flex_margin[f];
+    mjtNum margin = override_margin ? override_margin : m->flex_margin[f];
     mjtNum bound = m->flex_radius[f] + margin;
     aamm[0] -= bound;
     aamm[1] -= bound;
diff --git a/src/engine/engine_inline.h b/src/engine/engine_inline.h
index bcf07756..dbe4b8f3 100644
--- a/src/engine/engine_inline.h
+++ b/src/engine/engine_inline.h
@@ -189,6 +189,21 @@ void mji_mulMatTMat3(mjtNum* restrict res, const mjtNum mat1[9], const mjtNum ma
 }
 
 
+// transpose 3x3 matrix
+static inline
+void mji_transpose3(mjtNum* restrict res, const mjtNum mat[9]) {
+  res[0] = mat[0];
+  res[1] = mat[3];
+  res[2] = mat[6];
+  res[3] = mat[1];
+  res[4] = mat[4];
+  res[5] = mat[7];
+  res[6] = mat[2];
+  res[7] = mat[5];
+  res[8] = mat[8];
+}
+
+
 //------------------------------ 4D vector and matrix-vector operations ----------------------------
 
 // res = vec
diff --git a/src/engine/engine_util_misc.c b/src/engine/engine_util_misc.c
index 8a867b1d..abaf7055 100644
--- a/src/engine/engine_util_misc.c
+++ b/src/engine/engine_util_misc.c
@@ -1304,21 +1304,13 @@ void mju_printMatSparse(const mjtNum* mat, int nr,
 
 // min function, avoid re-evaluation
 mjtNum mju_min(mjtNum a, mjtNum b) {
-  if (a <= b) {
-    return a;
-  } else {
-    return b;
-  }
+  return a <= b ? a : b;
 }
 
 
 // max function, avoid re-evaluation
 mjtNum mju_max(mjtNum a, mjtNum b) {
-  if (a >= b) {
-    return a;
-  } else {
-    return b;
-  }
+  return a >= b ? a : b;
 }