dm_control: Import of refs/pull/468/head

PiperOrigin-RevId: 631836013 Change-Id: I4544185aa0bbd501d9e81db64d253cd2da8382d4
google-deepmind · May 8, 2024 · 19f7f68 · 19f7f68
1 parent db4a392
commit 19f7f68
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 4 deletions.
diff --git a/dm_control/mjcf/schema.py b/dm_control/mjcf/schema.py
@@ -160,8 +160,10 @@ def _parse_attribute(attribute_xml):
   else:
     try:
       attribute_callable = _SCALAR_TYPE_MAP[attribute_type]
-    except KeyError:
-      raise ValueError('Invalid attribute type: {}'.format(attribute_type))
+    except KeyError as exc:
+      raise ValueError(
+          'Invalid attribute type: {}'.format(attribute_type)
+      ) from exc
 
   return AttributeSpec(
       name=name, type=attribute_callable, required=required,
@@ -215,7 +217,7 @@ def _attachment_frame_spec(is_world_attachment):
   body_spec = MUJOCO.children['worldbody'].children['body']
   # 'name' and 'childclass' attributes are excluded.
   for attrib_name in (
-      'mocap', 'pos', 'quat', 'axisangle', 'xyaxes', 'zaxis', 'euler'):
+      'mocap', 'pos', 'quat', 'axisangle', 'xyaxes', 'zaxis', 'euler', 'user'):
     frame_spec.attributes[attrib_name] = copy.deepcopy(
         body_spec.attributes[attrib_name])
 
@@ -258,4 +260,3 @@ def override_schema(schema_xml_path):
   MUJOCO = parse_schema(schema_xml_path)
   FINDABLE_NAMESPACES = frozenset(
       collect_namespaces(MUJOCO).union(_ADDITIONAL_FINDABLE_NAMESPACES))
-
diff --git a/dm_control/mjcf/schema.xml b/dm_control/mjcf/schema.xml
@@ -1308,6 +1308,7 @@
             <attribute name="zaxis" type="array" array_type="float" array_size="3"/>
             <attribute name="euler" type="array" array_type="float" array_size="3"/>
             <attribute name="gravcomp" type="float"/>
+            <attribute name="user" type="array" array_type="float"/>
           </attributes>
           <children>
             <element name="plugin" repeated="true" namespace="body">