Add insert method to mjcf.element.

PiperOrigin-RevId: 469135739
Change-Id: I60d239051e8bbded0f90add424418f7bd5c63c2f

Author: yuvaltassa

dm_control/mjcf/element.py

@@ -592,6 +592,23 @@ def add(self, element_name, **kwargs):
       ValueError: If the 'element_name' is not a valid child, or if an invalid
         attribute is specified in `kwargs`.
 
+    Returns:
+      An `mjcf.Element` corresponding to the newly created child element.
+    """
+    return self.insert(element_name, position=None, **kwargs)
+
+  def insert(self, element_name, position, **kwargs):
+    """Add a new child element to this element.
+
+    Args:
+      element_name: The tag of the element to add.
+      position: Where to insert the new element.
+      **kwargs: Attributes of the new element being created.
+
+    Raises:
+      ValueError: If the 'element_name' is not a valid child, or if an invalid
+        attribute is specified in `kwargs`.
+
     Returns:
       An `mjcf.Element` corresponding to the newly created child element.
     """
@@ -604,7 +621,10 @@ def add(self, element_name, **kwargs):
       raise ValueError('A <{}> child already exists, please access it directly.'
                        .format(element_name))
     new_element = _make_element(child_spec, self, attributes=kwargs)
-    self._children.append(new_element)
+    if position is not None:
+      self._children.insert(position, new_element)
+    else:
+      self._children.append(new_element)
     self.namescope.increment_revision()
     return new_element
 

dm_control/mjcf/element_test.py

@@ -198,6 +198,26 @@ def testAdd(self):
       mujoco.add('default')
     self.assertIsNotNone(mujoco.default)
 
+  def testInsert(self):
+    mujoco = element.RootElement(model='test')
+
+    # add in order
+    mujoco.worldbody.add('body', name='0')
+    mujoco.worldbody.add('body', name='1')
+    mujoco.worldbody.add('body', name='2')
+
+    # insert into position 0, check order
+    mujoco.worldbody.insert('body', name='foo', position=0)
+    expected_order = ['foo', '0', '1', '2']
+    for i, child in enumerate(mujoco.worldbody._children):
+      self.assertEqual(child.name, expected_order[i])
+
+    # insert into position -1, check order
+    mujoco.worldbody.insert('body', name='bar', position=-1)
+    expected_order = ['foo', '0', '1', 'bar', '2']
+    for i, child in enumerate(mujoco.worldbody._children):
+      self.assertEqual(child.name, expected_order[i])
+
   def testAddWithInvalidAttribute(self):
     mujoco = element.RootElement(model='test')
     with self.assertRaisesRegex(AttributeError, 'not a valid attribute'):