From cf6591de271f64d04015fa77cdfafe87a2496e7f Mon Sep 17 00:00:00 2001 From: Shubh Date: Sat, 15 Aug 2026 20:08:16 +0530 Subject: [PATCH] Support include below the root element. MuJoCo allows wherever a child element is allowed and puts the contents of the named file in its place, but the parser only looked for it as a child of . Anything deeper was left in the tree and reached the schema, which reported it as a KeyError on the tag rather than as anything to do with includes. Expand those in place before parsing, recursively, so a file included from a subdirectory can include further files relative to itself. Top level includes still go through include_copy as whole models so they keep their own asset directories. Fixes #529. Signed-off-by: Shubh --- dm_control/mjcf/element_test.py | 25 +++++++-- dm_control/mjcf/parser.py | 53 +++++++++++++++++++ dm_control/mjcf/test_assets/included_body.xml | 6 +++ .../mjcf/test_assets/included_defaults.xml | 5 ++ .../test_assets/model_with_nested_include.xml | 11 ++++ 5 files changed, 97 insertions(+), 3 deletions(-) create mode 100644 dm_control/mjcf/test_assets/included_body.xml create mode 100644 dm_control/mjcf/test_assets/included_defaults.xml create mode 100644 dm_control/mjcf/test_assets/model_with_nested_include.xml diff --git a/dm_control/mjcf/element_test.py b/dm_control/mjcf/element_test.py index a8deace1..8af23f94 100644 --- a/dm_control/mjcf/element_test.py +++ b/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,7 +862,8 @@ 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: @@ -867,10 +871,25 @@ def testParseFromFile(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 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 , 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() diff --git a/dm_control/mjcf/parser.py b/dm_control/mjcf/parser.py index 681f0293..eda2e358 100644 --- a/dm_control/mjcf/parser.py +++ b/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 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 `` 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 : got <{}> in ' + '{!r}'.format(included_root.tag, filename)) + return included_root + + +def _expand_nested_includes(xml_element, model_dir, assets): + """Replaces `` tags below the root with the contents of the file. + + MuJoCo allows `` 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. diff --git a/dm_control/mjcf/test_assets/included_body.xml b/dm_control/mjcf/test_assets/included_body.xml new file mode 100644 index 00000000..6255d050 --- /dev/null +++ b/dm_control/mjcf/test_assets/included_body.xml @@ -0,0 +1,6 @@ + + + + + + diff --git a/dm_control/mjcf/test_assets/included_defaults.xml b/dm_control/mjcf/test_assets/included_defaults.xml new file mode 100644 index 00000000..60271bd9 --- /dev/null +++ b/dm_control/mjcf/test_assets/included_defaults.xml @@ -0,0 +1,5 @@ + + + + + diff --git a/dm_control/mjcf/test_assets/model_with_nested_include.xml b/dm_control/mjcf/test_assets/model_with_nested_include.xml new file mode 100644 index 00000000..f88175bc --- /dev/null +++ b/dm_control/mjcf/test_assets/model_with_nested_include.xml @@ -0,0 +1,11 @@ + + + + + + + + + + +