diff --git a/python/mujoco/bindings_test.py b/python/mujoco/bindings_test.py
index 21e048ab..d7cb142f 100644
--- a/python/mujoco/bindings_test.py
+++ b/python/mujoco/bindings_test.py
@@ -699,6 +699,49 @@ class MuJoCoBindingsTest(parameterized.TestCase):
self.assertEmpty(self.data.contact)
self.assertEmpty(self.data.efc_id)
+ def test_realloc_island(self):
+ # Test allocation on fresh data (on its own)
+ nisland = 2
+ nidof = 4
+ mujoco._functions._realloc_island(self.data, nisland=nisland, nidof=nidof)
+ self.assertEqual(self.data.nisland, nisland)
+ self.assertEqual(self.data.nidof, nidof)
+ self.assertEqual(self.data.island_nv.shape, (nisland,))
+ self.assertEqual(self.data.ifrc_smooth.shape, (nidof,))
+
+ # Test allocation after _realloc_con_efc
+ nefc = 10
+ mujoco._functions._realloc_con_efc(self.data, ncon=0, nefc=nefc)
+
+ nisland = 3
+ nidof = 5
+ mujoco._functions._realloc_island(self.data, nisland=nisland, nidof=nidof)
+
+ self.assertEqual(self.data.nisland, nisland)
+ self.assertEqual(self.data.nidof, nidof)
+ self.assertEqual(self.data.island_nv.shape, (nisland,))
+ self.assertEqual(self.data.ifrc_smooth.shape, (nidof,))
+
+ # Test re-allocation (calling it again with different sizes)
+ nisland2 = 4
+ nidof2 = 6
+ mujoco._functions._realloc_island(self.data, nisland=nisland2, nidof=nidof2)
+
+ self.assertEqual(self.data.nisland, nisland2)
+ self.assertEqual(self.data.nidof, nidof2)
+ self.assertEqual(self.data.island_nv.shape, (nisland2,))
+ self.assertEqual(self.data.ifrc_smooth.shape, (nidof2,))
+
+ # Test insufficient memory handling
+ expected_error = (
+ r'Insufficient arena memory, currently allocated memory=' +
+ r'"[0-9]+[A-Z]?". Increase using .'
+ )
+ with self.assertRaisesRegex(mujoco.FatalError, expected_error):
+ mujoco._functions._realloc_island(self.data, 100000000, 100000000)
+ self.assertEqual(self.data.nisland, 0)
+ self.assertEqual(self.data.nidof, 0)
+
def test_mj_struct_list_equality(self):
model2 = mujoco.MjModel.from_xml_string(TEST_XML)
data2 = mujoco.MjData(model2)
diff --git a/python/mujoco/functions.cc b/python/mujoco/functions.cc
index e56643b0..b8673a0e 100644
--- a/python/mujoco/functions.cc
+++ b/python/mujoco/functions.cc
@@ -1696,6 +1696,7 @@ PYBIND11_MODULE(_functions, pymodule, pybind11::mod_gil_not_used()) {
#define X(type, name, nr, nc) data->name = nullptr;
MJDATA_ARENA_POINTERS_SOLVER
MJDATA_ARENA_POINTERS_DUAL
+ MJDATA_ARENA_POINTERS_ISLAND
#undef X
};
@@ -1744,6 +1745,74 @@ PYBIND11_MODULE(_functions, pymodule, pybind11::mod_gil_not_used()) {
},
py::arg("d"), py::arg("ncon"), py::arg("nefc"), py::arg("nJ") = -1,
py::call_guard());
+
+ pymodule.def(
+ "_realloc_island",
+ [](MjDataWrapper& d, int nisland, int nidof) {
+ raw::MjData* data = d.get();
+
+ size_t parena_start = data->parena;
+ // Find island block start in arena to reclaim memory on re-allocation.
+ char* min_ptr = nullptr;
+#define X(type, name, nr, nc) \
+ if (data->name && \
+ (!min_ptr || reinterpret_cast(data->name) < min_ptr)) { \
+ min_ptr = reinterpret_cast(data->name); \
+ }
+ MJDATA_ARENA_POINTERS_ISLAND
+#undef X
+ if (min_ptr && data->arena) {
+ parena_start = min_ptr - static_cast(data->arena);
+ }
+
+ auto cleanup = [](raw::MjData* data, size_t target_parena) {
+#define X(type, name, nr, nc) data->name = nullptr;
+ MJDATA_ARENA_POINTERS_ISLAND
+#undef X
+ data->nisland = 0;
+ data->nidof = 0;
+ data->parena = target_parena;
+#ifdef ADDRESS_SANITIZER
+ ASAN_POISON_MEMORY_REGION(
+ static_cast(data->arena) + target_parena,
+ data->narena - data->pstack - target_parena);
+#endif
+ };
+
+ cleanup(data, parena_start);
+
+ char error_msg[128];
+ error_msg[0] = '\0';
+ const char* error_msg_fmt =
+ "Insufficient arena memory, currently allocated memory=\"%s\". "
+ "Increase using .";
+
+ data->nisland = nisland;
+ data->nidof = nidof;
+
+#undef MJ_M
+#define MJ_M(x) d.model().get()->x
+#undef MJ_D
+#define MJ_D(x) data->x
+#define X(type, name, nr, nc) \
+ data->name = static_cast(InterceptMjErrors(::mj_arenaAllocByte)( \
+ data, sizeof(type) * (nr) * (nc), alignof(type))); \
+ if (!data->name) { \
+ cleanup(data, parena_start); \
+ std::snprintf(error_msg, sizeof(error_msg), error_msg_fmt, \
+ mju_writeNumBytes(data->narena)); \
+ throw FatalError(error_msg); \
+ }
+
+ MJDATA_ARENA_POINTERS_ISLAND
+#undef X
+#undef MJ_D
+#define MJ_D(x) x
+#undef MJ_M
+#define MJ_M(x) x
+ },
+ py::arg("d"), py::arg("nisland"), py::arg("nidof"),
+ py::call_guard());
} // PYBIND11_MODULE NOLINT(readability/fn_size)
} // namespace
} // namespace mujoco::python