From 648a03c2cf136f67ceb6ab9f64e2d80bbf901d40 Mon Sep 17 00:00:00 2001
From: Alessio Quaglino <quaglino@google.com>
Date: Tue, 1 Apr 2025 05:20:59 -0700
Subject: [PATCH] Override inertia of body with visual geoms if discardvisual
 is true.

Fixes #2546.

PiperOrigin-RevId: 742650702
Change-Id: I753c1e1c89da732ece100676644bbecf9d6373bb
---
 src/user/user_model.cc     | 31 ++++++++++++++++++++++++++-----
 src/user/user_objects.cc   | 10 ----------
 test/user/user_api_test.cc |  3 +--
 3 files changed, 27 insertions(+), 17 deletions(-)

diff --git a/src/user/user_model.cc b/src/user/user_model.cc
index a63f13c9..11ef0755 100644
--- a/src/user/user_model.cc
+++ b/src/user/user_model.cc
@@ -558,7 +558,7 @@ void mjCModel::RemoveFromList(std::vector<T*>& list, const mjCModel& other) {
 template <>
 void mjCModel::DeleteAll<mjCKey>(std::vector<mjCKey*>& elements) {
   for (mjCKey* element : elements) {
-    delete element;
+    element->Release();
   }
   elements.clear();
 }
@@ -1553,7 +1553,7 @@ static void DeleteElements(std::vector<T*>& elements,
   int i = 0;
   for (int j=0; j < elements.size(); j++) {
     if (discard[j]) {
-      delete elements[j];
+      elements[j]->Release();
     } else {
       elements[i] = elements[j];
       i++;
@@ -1611,7 +1611,7 @@ void mjCModel::DeleteAll<mjCMaterial>(std::vector<mjCMaterial*>& elements) {
   DeleteMaterial(sites_);
   DeleteMaterial(tendons_);
   for (mjCMaterial* element : elements) {
-    delete element;
+    element->Release();
   }
   elements.clear();
 }
@@ -1621,7 +1621,7 @@ template <>
 void mjCModel::DeleteAll<mjCTexture>(std::vector<mjCTexture*>& elements) {
   DeleteAllTextures(materials_);
   for (mjCTexture* element : elements) {
-    delete element;
+    element->Release();
   }
   elements.clear();
 }
@@ -1781,7 +1781,6 @@ void mjCModel::IndexAssets(bool discard) {
     }
   }
 
-  // discard visual meshes and geoms
   if (discard) {
     std::vector<bool> discard_mesh(meshes_.size(), false);
     std::vector<bool> discard_geom(geoms_.size(), false);
@@ -1795,6 +1794,28 @@ void mjCModel::IndexAssets(bool discard) {
       return geom->IsVisual();
     });
 
+    // update inertia in bodies
+    for (auto body : bodies_) {
+      if (body->spec.explicitinertial) {
+        continue;
+      }
+      for (auto geom : body->geoms) {
+        if (geom->IsVisual()) {
+          if (compiler.inertiafromgeom == mjINERTIAFROMGEOM_TRUE) {
+            compiler.inertiafromgeom = mjINERTIAFROMGEOM_AUTO;
+          }
+          body->explicitinertial = true;  // for XML writer
+          body->spec.explicitinertial = true;
+          body->spec.mass = body->mass;
+          mjuu_copyvec(body->spec.ipos, body->ipos, 3);
+          mjuu_copyvec(body->spec.iquat, body->iquat, 4);
+          mjuu_copyvec(body->spec.inertia, body->inertia, 3);
+          break;
+        }
+      }
+    }
+
+    // discard visual meshes and geoms
     Delete(meshes_, discard_mesh);
     Delete(geoms_, discard_geom);
   }
diff --git a/src/user/user_objects.cc b/src/user/user_objects.cc
index d65ac68b..d46c51df 100644
--- a/src/user/user_objects.cc
+++ b/src/user/user_objects.cc
@@ -1976,16 +1976,6 @@ void mjCBody::Compile(void) {
     }
   }
 
-  // if discarding visual geoms, use explicit inertias
-  if (compiler->discardvisual) {
-    for (int j=0; j < geoms.size(); j++) {
-      if (geoms[j]->IsVisual()) {
-        explicitinertial = true;
-        break;
-      }
-    }
-  }
-
   // free joint alignment, phase 2 (transform sites, cameras and lights)
   if (align_free) {
     // frames have already been compiled and applied to children
diff --git a/test/user/user_api_test.cc b/test/user/user_api_test.cc
index 3837c423..d07f0388 100644
--- a/test/user/user_api_test.cc
+++ b/test/user/user_api_test.cc
@@ -520,8 +520,7 @@ TEST_F(PluginTest, RecompileCompare) {
         // if file is meant to fail, skip it
         if (absl::StrContains(p.path().string(), "malformed_") ||
             absl::StrContains(p.path().string(), "touch_grid") ||
-            absl::StrContains(p.path().string(), "cow") ||
-            absl::StrContains(p.path().string(), "discardvisual")) {
+            absl::StrContains(p.path().string(), "cow")) {
           continue;
         }