From ad49733db2772dac88eb8be81936c08c639f2735 Mon Sep 17 00:00:00 2001
From: Yuval Tassa <tassa@google.com>
Date: Sat, 1 Nov 2025 05:53:56 -0700
Subject: [PATCH] Improve error message for insufficient arena memory in
 `_realloc_con_efc`.

PiperOrigin-RevId: 826845370
Change-Id: I09467dd44a6f9f01b89ff6c1d74cde493e84eb1c
---
 python/mujoco/bindings_test.py |  7 +++++--
 python/mujoco/functions.cc     | 35 ++++++++++++++++++++--------------
 2 files changed, 26 insertions(+), 16 deletions(-)

diff --git a/python/mujoco/bindings_test.py b/python/mujoco/bindings_test.py
index 90ac1302..2d1f9e2f 100644
--- a/python/mujoco/bindings_test.py
+++ b/python/mujoco/bindings_test.py
@@ -632,8 +632,11 @@ class MuJoCoBindingsTest(parameterized.TestCase):
     self.assertEqual(self.data.efc_J.shape, (nj,))
     self.assertEqual(self.data.efc_KBIP.shape, (nefc, 4))
 
-    expected_error = 'insufficient arena memory available'
-    with self.assertRaisesWithLiteralMatch(mujoco.FatalError, expected_error):
+    expected_error = (
+        r'Insufficient arena memory, currently allocated memory=' +
+        r'"[0-9]+[A-Z]?". Increase using <size memory="X"/>.'
+    )
+    with self.assertRaisesRegex(mujoco.FatalError, expected_error):
       mujoco._functions._realloc_con_efc(self.data, 100000000, 100000000)
     self.assertEmpty(self.data.contact)
     self.assertEmpty(self.data.efc_id)
diff --git a/python/mujoco/functions.cc b/python/mujoco/functions.cc
index 78bf94b3..27e5b173 100644
--- a/python/mujoco/functions.cc
+++ b/python/mujoco/functions.cc
@@ -16,6 +16,7 @@
 
 #include <array>
 #include <cstdint>
+#include <cstdio>
 #include <memory>
 #include <string>
 #include <optional>
@@ -264,27 +265,24 @@ PYBIND11_MODULE(_functions, pymodule) {
             m, d, x.data(), y.data(), y.rows());
       });
   DEF_WITH_OMITTED_PY_ARGS(traits::mj_solveM2, "n")(
-      pymodule,
-      [](const raw::MjModel* m, raw::MjData* d, Eigen::Ref<EigenArrayXX> x,
-         Eigen::Ref<const EigenArrayXX> y, Eigen::Ref<const EigenArrayXX> sqrtInvD) {
+      pymodule, [](const raw::MjModel* m, raw::MjData* d,
+                   Eigen::Ref<EigenArrayXX> x, Eigen::Ref<const EigenArrayXX> y,
+                   Eigen::Ref<const EigenArrayXX> sqrtInvD) {
         if (x.rows() != y.rows()) {
           throw py::type_error(
               "the first dimension of x and y should be of the same size");
         }
         if (x.cols() != m->nv) {
-          throw py::type_error(
-              "the last dimension of x should be of size nv");
+          throw py::type_error("the last dimension of x should be of size nv");
         }
         if (y.cols() != m->nv) {
-          throw py::type_error(
-              "the last dimension of y should be of size nv");
+          throw py::type_error("the last dimension of y should be of size nv");
         }
         if (sqrtInvD.size() != m->nv) {
-          throw py::type_error(
-              "the size of sqrtInvD should be nv");
+          throw py::type_error("the size of sqrtInvD should be nv");
         }
-        return InterceptMjErrors(::mj_solveM2)(
-            m, d, x.data(), y.data(), sqrtInvD.data(), y.rows());
+        return InterceptMjErrors(::mj_solveM2)(m, d, x.data(), y.data(),
+                                               sqrtInvD.data(), y.rows());
       });
   Def<traits::mj_comVel>(pymodule);
   Def<traits::mj_passive>(pymodule);
@@ -1533,6 +1531,11 @@ PYBIND11_MODULE(_functions, pymodule) {
 #undef X
         };
 
+        char error_msg[128];
+        error_msg[0] = '\0';
+        const char* error_msg_fmt =
+            "Insufficient arena memory, currently allocated memory=\"%s\". "
+            "Increase using <size memory=\"X\"/>.";
         cleanup(data, nJ);
         data->ncon = ncon;
         data->nefc = nefc;
@@ -1542,7 +1545,9 @@ PYBIND11_MODULE(_functions, pymodule) {
                 data, ncon * sizeof(raw::MjContact), alignof(raw::MjContact)));
         if (!data->contact) {
           cleanup(data, nJ);
-          throw FatalError("insufficient arena memory available");
+          std::snprintf(error_msg, sizeof(error_msg), error_msg_fmt,
+                        mju_writeNumBytes(data->narena));
+          throw FatalError(error_msg);
         }
 
 #undef MJ_M
@@ -1553,8 +1558,10 @@ PYBIND11_MODULE(_functions, pymodule) {
   data->name = static_cast<type*>(InterceptMjErrors(::mj_arenaAllocByte)( \
       data, sizeof(type) * (nr) * (nc), alignof(type)));                  \
   if (!data->name) {                                                      \
-    cleanup(data, nJ);                                                        \
-    throw FatalError("insufficient arena memory available");              \
+    cleanup(data, nJ);                                                    \
+    std::snprintf(error_msg, sizeof(error_msg), error_msg_fmt,            \
+                  mju_writeNumBytes(data->narena));                       \
+    throw FatalError(error_msg);                                          \
   }
 
         MJDATA_ARENA_POINTERS_SOLVER