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

Allow _realloc_con_efc to also allocate new nJ arrays.

Browse Source 
PiperOrigin-RevId: 761252071
Change-Id: Idce822c6f2b8f2d81a5d43228f5dc35281c5ad07
This commit is contained in:
Erik Frey
2025-05-20 15:25:10 -07:00
committed by Copybara-Service
parent 5051520733
commit 49c33716a1
2 changed files with 11 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/mujoco/bindings_test.py
+3 -1
View File
@@ -623,10 +623,12 @@ class MuJoCoBindingsTest(parameterized.TestCase):
ncon = 13 ncon = 13
nefc = 17 nefc = 17
mujoco._functions._realloc_con_efc(self.data, ncon=ncon, nefc=nefc) nj = 21
mujoco._functions._realloc_con_efc(self.data, ncon=ncon, nefc=nefc, nJ=nj)
self.assertLen(self.data.contact, ncon) self.assertLen(self.data.contact, ncon)
self.assertEqual(self.data.efc_id.shape, (nefc,)) self.assertEqual(self.data.efc_id.shape, (nefc,))
self.assertEqual(self.data.efc_J.shape, (nj,))
self.assertEqual(self.data.efc_KBIP.shape, (nefc, 4)) self.assertEqual(self.data.efc_KBIP.shape, (nefc, 4))
expected_error = 'insufficient arena memory available' expected_error = 'insufficient arena memory available'
python/mujoco/functions.cc
+8 -6
View File
@@ -1467,10 +1467,10 @@ PYBIND11_MODULE(_functions, pymodule) {
pymodule.def( pymodule.def(
"_realloc_con_efc", "_realloc_con_efc",
[](MjDataWrapper& d, int ncon, int nefc) { [](MjDataWrapper& d, int ncon, int nefc, int nJ) {
raw::MjData* data = d.get(); raw::MjData* data = d.get();
auto cleanup = [](raw::MjData* data) { auto cleanup = [](raw::MjData* data, int nJ) {
#ifdef ADDRESS_SANITIZER #ifdef ADDRESS_SANITIZER
ASAN_POISON_MEMORY_REGION( ASAN_POISON_MEMORY_REGION(
static_cast<char*>(data->arena), static_cast<char*>(data->arena),
@@ -1479,6 +1479,7 @@ PYBIND11_MODULE(_functions, pymodule) {
data->parena = 0; data->parena = 0;
data->ncon = 0; data->ncon = 0;
data->nefc = 0; data->nefc = 0;
if (nJ > -1) data->nJ = 0;
data->contact = static_cast<raw::MjContact*>(data->arena); data->contact = static_cast<raw::MjContact*>(data->arena);
#define X(type, name, nr, nc) data->name = nullptr; #define X(type, name, nr, nc) data->name = nullptr;
MJDATA_ARENA_POINTERS_SOLVER MJDATA_ARENA_POINTERS_SOLVER
@@ -1486,14 +1487,15 @@ PYBIND11_MODULE(_functions, pymodule) {
#undef X #undef X
}; };
cleanup(data); cleanup(data, nJ);
data->ncon = ncon; data->ncon = ncon;
data->nefc = nefc; data->nefc = nefc;
if (nJ > -1) data->nJ = nJ;
data->contact = data->contact =
static_cast<raw::MjContact*>(InterceptMjErrors(::mj_arenaAllocByte)( static_cast<raw::MjContact*>(InterceptMjErrors(::mj_arenaAllocByte)(
data, ncon * sizeof(raw::MjContact), alignof(raw::MjContact))); data, ncon * sizeof(raw::MjContact), alignof(raw::MjContact)));
if (!data->contact) { if (!data->contact) {
cleanup(data); cleanup(data, nJ);
throw FatalError("insufficient arena memory available"); throw FatalError("insufficient arena memory available");
} }
@@ -1505,7 +1507,7 @@ PYBIND11_MODULE(_functions, pymodule) {
data->name = static_cast<type*>(InterceptMjErrors(::mj_arenaAllocByte)( \ data->name = static_cast<type*>(InterceptMjErrors(::mj_arenaAllocByte)( \
data, sizeof(type) * (nr) * (nc), alignof(type))); \ data, sizeof(type) * (nr) * (nc), alignof(type))); \
if (!data->name) { \ if (!data->name) { \
cleanup(data); \ cleanup(data, nJ); \
throw FatalError("insufficient arena memory available"); \ throw FatalError("insufficient arena memory available"); \
} }
@@ -1519,7 +1521,7 @@ PYBIND11_MODULE(_functions, pymodule) {
#undef MJ_M #undef MJ_M
#define MJ_M(x) x #define MJ_M(x) x
}, },
py::arg("d"), py::arg("ncon"), py::arg("nefc"), py::arg("d"), py::arg("ncon"), py::arg("nefc"), py::arg("nJ") = -1,
py::call_guard<py::gil_scoped_release>()); py::call_guard<py::gil_scoped_release>());
} // PYBIND11_MODULE NOLINT(readability/fn_size) } // PYBIND11_MODULE NOLINT(readability/fn_size)
} // namespace } // namespace