From 1c6609c2a9af28cb2197b9217a45cddcf9a89952 Mon Sep 17 00:00:00 2001 From: Saran Tunyasuvunakool Date: Tue, 12 Apr 2022 12:49:17 -0700 Subject: [PATCH] Don't return ndarray with empty shape through named indexers. NumPy arrays with an empty shape cannot be assigned via the usual syntax. Fixes: #238. Related: #237. PiperOrigin-RevId: 441265437 Change-Id: Ib0de22c83f9babe9a97682dc5a32660c198ad68c --- python/mujoco/bindings_test.py | 21 ++++++++++++++++++++- python/mujoco/indexers.cc | 5 +++++ 2 files changed, 25 insertions(+), 1 deletion(-) diff --git a/python/mujoco/bindings_test.py b/python/mujoco/bindings_test.py index 65165953c03..71bd3b3327f 100644 --- a/python/mujoco/bindings_test.py +++ b/python/mujoco/bindings_test.py @@ -41,7 +41,7 @@ - + @@ -51,6 +51,9 @@ + + + """ @@ -144,6 +147,22 @@ def test_array_keeps_struct_alive(self): del qpos_spring self.assertEqual(sys.getrefcount(capsule) - base_refcount, 1) + def test_named_indexing_actuator_ctrl(self): + actuator_id = mujoco.mj_name2id( + self.model, mujoco.mjtObj.mjOBJ_ACTUATOR, 'myactuator') + self.assertIs(self.data.actuator('myactuator'), + self.data.actuator(actuator_id)) + self.assertIs(self.data.actuator('myactuator').ctrl, + self.data.actuator(actuator_id).ctrl) + self.assertEqual(self.data.actuator('myactuator').ctrl.shape, (1,)) + + # Test that the indexer is returning a view into the underlying struct. + ctrl_from_indexer = self.data.actuator('myactuator').ctrl + self.data.ctrl[actuator_id] = 5 + np.testing.assert_array_equal(ctrl_from_indexer, [5]) + self.data.actuator('myactuator').ctrl = 7 + np.testing.assert_array_equal(self.data.ctrl[actuator_id], [7]) + def test_named_indexing_geom_size(self): box_id = mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_GEOM, 'mybox') self.assertIs(self.model.geom('mybox'), self.model.geom(box_id)) diff --git a/python/mujoco/indexers.cc b/python/mujoco/indexers.cc index 01cb1cf5ced..0e688afb287 100644 --- a/python/mujoco/indexers.cc +++ b/python/mujoco/indexers.cc @@ -102,6 +102,11 @@ py::array_t MakeArray(T* base_ptr, int index, std::vector&& shape, offset = m.tuple_adr[index]; shape.insert(shape.begin(), m.tuple_size[index]); } else { + // Do not return a NumPy array with shape () since these aren't very nice + // to work with. Instead, always return singleton arrays with shape (1,). + if (shape.empty()) { + shape.push_back(1); + } int size = 1; for (int s : shape) { size *= s;