Don't return ndarray with empty shape through named indexers.

NumPy arrays with an empty shape cannot be assigned via the usual syntax. Fixes: google-deepmind#238. Related: google-deepmind#237. PiperOrigin-RevId: 441265437 Change-Id: Ib0de22c83f9babe9a97682dc5a32660c198ad68c
kevinzakka · May 24, 2022 · 1c6609c · 1c6609c
1 parent 253bec2
commit 1c6609c
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 1 deletion.
diff --git a/python/mujoco/bindings_test.py b/python/mujoco/bindings_test.py
@@ -41,7 +41,7 @@
     <body>
       <inertial pos="0 0 0" mass="1" diaginertia="1 1 1"/>
       <site pos="0 0 -1" name="mysite" type="sphere"/>
-      <joint type="hinge" axis="0 1 0"/>
+      <joint name="myhinge" type="hinge" axis="0 1 0"/>
     </body>
     <body>
       <inertial pos="0 0 0" mass="1" diaginertia="1 1 1"/>
@@ -51,6 +51,9 @@
       <geom type="sphere" size="0.1"/>
     </body>
   </worldbody>
+  <actuator>
+    <position name="myactuator" joint="myhinge"/>
+  </actuator>
 </mujoco>
 """
 
@@ -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
@@ -102,6 +102,11 @@ py::array_t<T> MakeArray(T* base_ptr, int index, std::vector<int>&& 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;