chenlin/Mujoco_WASM
1
0
Fork 0
Code Issues Pull Requests Actions 2 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
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.assertEqual(self.data.efc_id.shape, (nefc,))
self.assertEqual(self.data.efc_J.shape, (nj,))
self.assertEqual(self.data.efc_KBIP.shape, (nefc, 4))
expected_error = 'insufficient arena memory available'
python/mujoco/functions.cc
+8 -6
View File
@@ -1467,10 +1467,10 @@ PYBIND11_MODULE(_functions, pymodule) {
pymodule.def(
"_realloc_con_efc",
[](MjDataWrapper& d, int ncon, int nefc) {
[](MjDataWrapper& d, int ncon, int nefc, int nJ) {
raw::MjData* data = d.get();
auto cleanup = [](raw::MjData* data) {
auto cleanup = [](raw::MjData* data, int nJ) {
#ifdef ADDRESS_SANITIZER
ASAN_POISON_MEMORY_REGION(
static_cast<char*>(data->arena),
@@ -1479,6 +1479,7 @@ PYBIND11_MODULE(_functions, pymodule) {
data->parena = 0;
data->ncon = 0;
data->nefc = 0;
if (nJ > -1) data->nJ = 0;
data->contact = static_cast<raw::MjContact*>(data->arena);
#define X(type, name, nr, nc) data->name = nullptr;
MJDATA_ARENA_POINTERS_SOLVER
@@ -1486,14 +1487,15 @@ PYBIND11_MODULE(_functions, pymodule) {
#undef X
};
cleanup(data);
cleanup(data, nJ);
data->ncon = ncon;
data->nefc = nefc;
if (nJ > -1) data->nJ = nJ;
data->contact =
static_cast<raw::MjContact*>(InterceptMjErrors(::mj_arenaAllocByte)(
data, ncon * sizeof(raw::MjContact), alignof(raw::MjContact)));
if (!data->contact) {
cleanup(data);
cleanup(data, nJ);
throw FatalError("insufficient arena memory available");
}
@@ -1505,7 +1507,7 @@ PYBIND11_MODULE(_functions, pymodule) {
data->name = static_cast<type*>(InterceptMjErrors(::mj_arenaAllocByte)( \
data, sizeof(type) * (nr) * (nc), alignof(type))); \
if (!data->name) { \
cleanup(data); \
cleanup(data, nJ); \
throw FatalError("insufficient arena memory available"); \
}
@@ -1519,7 +1521,7 @@ PYBIND11_MODULE(_functions, pymodule) {
#undef MJ_M
#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>());
} // PYBIND11_MODULE NOLINT(readability/fn_size)
} // namespace