Support include below the root element

dm_control/mjcf/element_test.py

@@ -39,6 +39,8 @@
 _TEXTURE_PATH = os.path.join(_ASSETS_DIR, 'textures/deepmind.png')
 _MESH_PATH = os.path.join(_ASSETS_DIR, 'meshes/cube.stl')
 _MODEL_WITH_INCLUDE_PATH = os.path.join(_ASSETS_DIR, 'model_with_include.xml')
+_MODEL_WITH_NESTED_INCLUDE_PATH = os.path.join(
+    _ASSETS_DIR, 'model_with_nested_include.xml')
 
 _MODEL_WITH_INVALID_FILENAMES = os.path.join(
     _ASSETS_DIR, 'model_with_invalid_filenames.xml')
@@ -850,7 +852,8 @@ def testResolveReferences(self):
 
   @parameterized.named_parameters(
       ('WithoutInclude', _TEST_MODEL_XML),
-      ('WithInclude', _MODEL_WITH_INCLUDE_PATH))
+      ('WithInclude', _MODEL_WITH_INCLUDE_PATH),
+      ('WithNestedInclude', _MODEL_WITH_NESTED_INCLUDE_PATH))
   def testParseFromString(self, model_path):
     with open(model_path) as xml_file:
       xml_string = xml_file.read()
@@ -859,18 +862,34 @@ def testParseFromString(self, model_path):
 
   @parameterized.named_parameters(
       ('WithoutInclude', _TEST_MODEL_XML),
-      ('WithInclude', _MODEL_WITH_INCLUDE_PATH))
+      ('WithInclude', _MODEL_WITH_INCLUDE_PATH),
+      ('WithNestedInclude', _MODEL_WITH_NESTED_INCLUDE_PATH))
   def testParseFromFile(self, model_path):
     model_dir, _ = os.path.split(model_path)
     with open(model_path) as xml_file:
       parser.from_file(xml_file, model_dir=model_dir)
 
   @parameterized.named_parameters(
       ('WithoutInclude', _TEST_MODEL_XML),
-      ('WithInclude', _MODEL_WITH_INCLUDE_PATH))
+      ('WithInclude', _MODEL_WITH_INCLUDE_PATH),
+      ('WithNestedInclude', _MODEL_WITH_NESTED_INCLUDE_PATH))
   def testParseFromPath(self, model_path):
     parser.from_path(model_path)
 
+  def testNestedIncludeIsPutInPlace(self):
+    mjcf_model = parser.from_path(_MODEL_WITH_NESTED_INCLUDE_PATH)
+
+    # The body was included inside <worldbody>, not merged at the root.
+    body = mjcf_model.find('body', 'included_body')
+    self.assertIsNotNone(body)
+    self.assertEqual(body.parent.tag, 'worldbody')
+
+    # The defaults were included inside the class that contained the tag.
+    nested = mjcf_model.default.find_all('default')[-1]
+    self.assertEqual(nested.dclass, 'nested')
+    self.assertEqual(nested.tendon.width, 0.002)
+    np.testing.assert_array_equal(nested.geom.rgba, [1, 0, 0, 1])
+
   def testGetAssetFromFile(self):
     with open(_TEXTURE_PATH, 'rb') as f:
       contents = f.read()

dm_control/mjcf/parser.py

@@ -216,6 +216,10 @@ def _parse(xml_root, escape_separators=False,
       # these are a schema violation.
       xml_root.remove(include_tag)
 
+    # MuJoCo also allows <include/> further down the tree, where the contents
+    # of the included file take the place of the tag.
+    _expand_nested_includes(xml_root, model_dir, assets)
+
     # Parse the main XML file.
     try:
       model = xml_root.attrib.pop('model')
@@ -237,6 +241,55 @@ def _parse(xml_root, escape_separators=False,
     return mjcf_root
 
 
+def _read_included_xml(include_tag, model_dir, assets):
+  """Returns the root `etree.Element` of the file an `<include/>` names."""
+  filename = include_tag.attrib['file']
+  try:
+    # First look for the included XML file in the assets dict.
+    contents = assets[filename]
+  except KeyError:
+    # If it's not present in the assets dict then load it from the filesystem.
+    contents = resources.GetResource(os.path.join(model_dir, filename))
+  included_root = etree.fromstring(contents)
+  if not included_root.tag.startswith('mujoco'):
+    raise ValueError(
+        'Root element of an included file should be <mujoco.*>: got <{}> in '
+        '{!r}'.format(included_root.tag, filename))
+  return included_root
+
+
+def _expand_nested_includes(xml_element, model_dir, assets):
+  """Replaces `<include/>` tags below the root with the contents of the file.
+
+  MuJoCo allows `<include/>` wherever a child element is allowed, not only at
+  the top level, and puts the children of the included file in its place. Top
+  level includes are handled separately in `_parse`, which merges them as whole
+  models so that they keep their own asset directories.
+
+  Args:
+    xml_element: The `etree.Element` whose descendants are to be expanded.
+    model_dir: Path to the directory containing the XML file being parsed.
+    assets: A dictionary of pre-loaded assets, of the form
+      `{filename: bytestring}`.
+  """
+  for child in list(xml_element):
+    if child.tag is etree.Comment or child.tag is etree.PI:
+      continue
+    if child.tag == 'include':
+      included_root = _read_included_xml(child, model_dir, assets)
+      # A file included from a subdirectory may include further files relative
+      # to itself.
+      included_dir = os.path.join(
+          model_dir, os.path.dirname(child.attrib['file']))
+      _expand_nested_includes(included_root, included_dir, assets)
+      index = xml_element.index(child)
+      xml_element.remove(child)
+      for offset, included_child in enumerate(list(included_root)):
+        xml_element.insert(index + offset, included_child)
+    else:
+      _expand_nested_includes(child, model_dir, assets)
+
+
 def _parse_children(xml_element, mjcf_element, escape_separators=False):
   """Parses all children of a given XML element into an MJCF element.
 

dm_control/mjcf/test_assets/included_body.xml

@@ -0,0 +1,6 @@
+<!-- Included by `model_with_nested_include.xml` -->
+<mujoco>
+  <body name="included_body">
+    <geom name="included_geom" class="nested" type="sphere" size="0.1"/>
+  </body>
+</mujoco>

dm_control/mjcf/test_assets/included_defaults.xml

@@ -0,0 +1,5 @@
+<!-- Included by `model_with_nested_include.xml` -->
+<mujoco>
+  <geom rgba="1 0 0 1"/>
+  <tendon width="0.002"/>
+</mujoco>

dm_control/mjcf/test_assets/model_with_nested_include.xml

@@ -0,0 +1,11 @@
+<!-- Model with an <include/> tag below the root element -->
+<mujoco>
+  <default>
+    <default class="nested">
+      <include file="included_defaults.xml"/>
+    </default>
+  </default>
+  <worldbody>
+    <include file="included_body.xml"/>
+  </worldbody>
+</mujoco>