From e973dc74dd8e324890fba2de6828e2a44ce5a60c Mon Sep 17 00:00:00 2001 From: Alessio Quaglino Date: Wed, 2 Oct 2024 11:13:45 -0700 Subject: [PATCH] Yield ownership of `char[]`, `vector`, and `vector` to `mjSpec` in Python bindings. PiperOrigin-RevId: 681522091 Change-Id: Icb3788d9ef6e914a307531e39ad8171c24f590af --- .../mujoco/codegen/generate_spec_bindings.py | 32 ++++------ python/mujoco/specs.cc | 36 +++++++++++ python/mujoco/specs_test.py | 59 +++++++++++++++++-- 3 files changed, 103 insertions(+), 24 deletions(-) diff --git a/python/mujoco/codegen/generate_spec_bindings.py b/python/mujoco/codegen/generate_spec_bindings.py index 9cdc7034..a7707bb3 100644 --- a/python/mujoco/codegen/generate_spec_bindings.py +++ b/python/mujoco/codegen/generate_spec_bindings.py @@ -93,15 +93,15 @@ def _array_binding_code( return f"""\ {classname}.def_property( "{varname}", - []({rawclassname}& self) -> py::array_t {{ - return py::array_t({field.extents[0]}, self.{fullvarname}); + []({rawclassname}& self) -> MjTypeVec {{ + return MjTypeVec(self.{fullvarname}, {field.extents[0]}); }}, []({rawclassname}& self, py::object rhs) {{ int i = 0; for (auto val : rhs) {{ self.{fullvarname}[i++] = py::cast(val); }} - }}, py::return_value_policy::reference_internal);""" + }}, py::return_value_policy::move);""" # all other array types return f"""\ {classname}.def_property( @@ -158,16 +158,13 @@ def _ptr_binding_code( self.{fullvarname}->push_back(py::cast<{vartype}>(val)); }} }}, py::return_value_policy::reference_internal);""" - elif vartype == 'mjByteVec': # C++ buffer -> Python list + elif vartype == 'mjByteVec': return f"""\ {classname}.def_property( "{varname}", - []({rawclassname}& self) -> py::list {{ - py::list list; - for (auto val : *self.{fullvarname}) {{ - list.append(val); - }} - return list; + []({rawclassname}& self) -> MjTypeVec {{ + return MjTypeVec(self.{fullvarname}->data(), + self.{fullvarname}->size()); }}, []({rawclassname}& self, py::object rhs) {{ self.{fullvarname}->clear(); @@ -175,17 +172,14 @@ def _ptr_binding_code( for (auto val : rhs) {{ self.{fullvarname}->push_back(py::cast(val)); }} - }}, py::return_value_policy::reference_internal);""" - elif vartype == 'mjStringVec': # C++ vector of strings -> Python list + }}, py::return_value_policy::move);""" + elif vartype == 'mjStringVec': return f"""\ {classname}.def_property( "{varname}", - []({rawclassname}& self) -> py::list {{ - py::list list; - for (auto val : *self.{fullvarname}) {{ - list.append(val); - }} - return list; + []({rawclassname}& self) -> MjTypeVec {{ + return MjTypeVec(self.{fullvarname}->data(), + self.{fullvarname}->size()); }}, []({rawclassname}& self, py::object rhs) {{ self.{fullvarname}->clear(); @@ -193,7 +187,7 @@ def _ptr_binding_code( for (auto val : rhs) {{ self.{fullvarname}->push_back(py::cast(val)); }} - }}, py::return_value_policy::reference_internal);""" + }}, py::return_value_policy::move);""" elif 'VecVec' in vartype: # C++ vector of vectors -> Python list of lists vartype = vartype.replace('mj', '').replace('VecVec', '').lower() return f"""\ diff --git a/python/mujoco/specs.cc b/python/mujoco/specs.cc index 8a1e7ecc..ad8f767a 100644 --- a/python/mujoco/specs.cc +++ b/python/mujoco/specs.cc @@ -20,6 +20,7 @@ #include #include // IWYU pragma: keep #include +#include #include // IWYU pragma: keep #include @@ -122,6 +123,38 @@ static raw::MjSpec* LoadSpecFileImpl( return spec; } +template +struct MjTypeVec { + MjTypeVec(T* data, int size) : ptr(data), size(size) {} + T* ptr; + int size; +}; + +template +void DefineArray(py::module& m, const std::string& typestr) { + using Class = MjTypeVec; + py::class_(m, typestr.c_str()) + .def(py::init([](T* data, int size) { return Class(data, size); })) + .def("__getitem__", + [](Class& v, int i) { + if (i < 0 || i >= v.size) { + throw py::index_error("Index out of range."); + } + return v.ptr[i]; + }) + .def("__setitem__", + [](Class& v, int i, T c) { + if (i < 0 || i >= v.size) { + throw py::index_error("Index out of range."); + } + v.ptr[i] = std::move(c); + }) + .def("__len__", [](Class& v) { return v.size; }) + .def("__iter__", [](Class& v) { + return py::make_iterator(v.ptr, v.ptr + v.size); + }, py::keep_alive<0, 1>(), py::return_value_policy::reference_internal); +}; + PYBIND11_MODULE(_specs, m) { auto structs_m = py::module::import("mujoco._structs"); py::function mjmodel_from_spec_ptr = @@ -161,6 +194,9 @@ PYBIND11_MODULE(_specs, m) { py::class_ mjOption(m, "MjOption"); py::class_ mjStatistic(m, "MjStatistic"); py::class_ mjVisual(m, "MjVisual"); + DefineArray(m, "MjCharVec"); + DefineArray(m, "MjStringVec"); + DefineArray(m, "MjByteVec"); // ============================= MJSPEC ===================================== mjSpec.def(py::init<>()); diff --git a/python/mujoco/specs_test.py b/python/mujoco/specs_test.py index 68e0f0cb..2bc0b599 100644 --- a/python/mujoco/specs_test.py +++ b/python/mujoco/specs_test.py @@ -34,9 +34,28 @@ class SpecsTest(absltest.TestCase): spec = mujoco.MjSpec() # Check that euler sequence order is set correctly. - self.assertEqual(spec.eulerseq[0], ord('x')) + self.assertEqual(spec.eulerseq[0], 'x') spec.eulerseq = ['z', 'y', 'x'] - self.assertEqual(spec.eulerseq[0], ord('z')) + self.assertEqual(spec.eulerseq[0], 'z') + + # Change single elements of euler sequence. + spec.eulerseq[0] = 'y' + spec.eulerseq[1] = 'z' + self.assertEqual(spec.eulerseq[0], 'y') + self.assertEqual(spec.eulerseq[1], 'z') + + # eulerseq is iterable + self.assertEqual('yzx', ''.join(spec.eulerseq)) + + # supports `len` + self.assertLen(spec.eulerseq, 3) + + # field checks for out-of-bound access on read and on write + with self.assertRaises(IndexError): + spec.eulerseq[3] = 'x' + + with self.assertRaises(IndexError): + spec.eulerseq[-1] = 'x' # Add a body, check that it has default orientation. body = spec.worldbody.add_body() @@ -50,10 +69,14 @@ class SpecsTest(absltest.TestCase): self.assertEqual(body.name, 'baz') # Change the position of the body and read it back. - body.pos = [1, 2, 3] - np.testing.assert_array_equal(body.pos, [1, 2, 3]) + body.pos = [4, 2, 3] + np.testing.assert_array_equal(body.pos, [4, 2, 3]) self.assertEqual(body.pos.shape, (3,)) + # Change single element of position. + body.pos[0] = 1 + np.testing.assert_array_equal(body.pos, [1, 2, 3]) + # Change the orientation of the body and read it back. body.quat = [0, 1, 0, 0] np.testing.assert_array_equal(body.quat, [0, 1, 0, 0]) @@ -147,7 +170,7 @@ class SpecsTest(absltest.TestCase): # Add tuple. tuple_ = spec.add_tuple(objprm=[2.0, 3.0, 5.0], objname=['obj']) np.testing.assert_array_equal(tuple_.objprm, [2.0, 3.0, 5.0]) - self.assertEqual(tuple_.objname, ['obj']) + self.assertEqual(tuple_.objname[0], 'obj') # Add flex. flex = spec.add_flex(friction=[1, 2, 3], texcoord=[1.0, 2.0, 3.0]) @@ -797,5 +820,31 @@ class SpecsTest(absltest.TestCase): self.assertEqual(model.stat.meansize, 0.06) self.assertEqual(model.vis.quality.shadowsize, 8192) + def test_assign_list_element(self): + spec = mujoco.MjSpec() + material = spec.add_material() + texture_index = mujoco.mjtTextureRole.mjTEXROLE_RGB + + # Assign a string to a list element. + material.textures[texture_index] = 'texture_name' + self.assertEqual(material.textures[texture_index], 'texture_name') + + # Assign a complete list + material.textures = ['', 'new_name', '', '', ''] + self.assertEqual(material.textures[texture_index], 'new_name') + + # textures is iterable + self.assertEqual('new_name', ''.join(material.textures)) + + # supports `len` + self.assertLen(material.textures, 5) + + # field checks for out-of-bound access on read and on write + with self.assertRaises(IndexError): + material.textures[5] = 'x' + + with self.assertRaises(IndexError): + material.textures[-1] = 'x' + if __name__ == '__main__': absltest.main()