# Copyright 2018 The dm_control Authors.

#

# Licensed under the Apache License, Version 2.0 (the "License");

# you may not use this file except in compliance with the License.

# You may obtain a copy of the License at

#

# http://www.apache.org/licenses/LICENSE-2.0

#

# Unless required by applicable law or agreed to in writing, software

# distributed under the License is distributed on an "AS IS" BASIS,

# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

# See the License for the specific language governing permissions and

# limitations under the License.

# ============================================================================

"""A Python object representation of Mujoco's MJCF schema.

The root schema is provided as a module-level constant `schema.MUJOCO`.

"""

import collections

import copy

import os

from dm_control.mjcf import attribute

from lxml import etree

from dm_control.utils import io as resources

_SCHEMA_XML_PATH = os.path.join(os.path.dirname(__file__), 'schema.xml')

_ARRAY_DTYPE_MAP = {

'int': int,

'float': float,

'string': str

}

_SCALAR_TYPE_MAP = {

'int': attribute.Integer,

'float': attribute.Float,

'string': attribute.String

}

ElementSpec = collections.namedtuple(

'ElementSpec', ('name', 'repeated', 'on_demand', 'identifier', 'namespace',

'attributes', 'children'))

AttributeSpec = collections.namedtuple(

'AttributeSpec', ('name', 'type', 'required',

'conflict_allowed', 'conflict_behavior', 'other_kwargs'))

# Additional namespaces that are not present in the MJCF schema but can

# be used in `find` and `find_all`.

_ADDITIONAL_FINDABLE_NAMESPACES = frozenset(['attachment_frame'])

def _str2bool(string):

"""Converts either 'true' or 'false' (not case-sensitively) into a boolean."""

if string is None:

return False

else:

string = string.lower()

if string == 'true':

return True

elif string == 'false':

return False

else:

raise ValueError(

'String should either be `true` or `false`: got {}'.format(string))

def parse_schema(schema_path):

"""Parses the schema XML.

Args:

schema_path: Path to the schema XML file.

Returns:

An `ElementSpec` for the root element in the schema.

"""

with resources.GetResourceAsFile(schema_path) as file_handle:

schema_xml = etree.parse(file_handle).getroot()

return _parse_element(schema_xml)

def _parse_element(element_xml):

"""Parses an <element> element in the schema."""

name = element_xml.get('name')

if not name:

raise ValueError('Element must always have a name')

repeated = _str2bool(element_xml.get('repeated'))

on_demand = _str2bool(element_xml.get('on_demand'))

attributes = collections.OrderedDict()

attributes_xml = element_xml.find('attributes')

if attributes_xml is not None:

for attribute_xml in attributes_xml.findall('attribute'):

attributes[attribute_xml.get('name')] = _parse_attribute(attribute_xml)

identifier = None

namespace = None

for attribute_spec in attributes.values():

if attribute_spec.type == attribute.Identifier:

identifier = attribute_spec.name

namespace = element_xml.get('namespace') or name

children = collections.OrderedDict()

children_xml = element_xml.find('children')

if children_xml is not None:

for child_xml in children_xml.findall('element'):

children[child_xml.get('name')] = _parse_element(child_xml)

element_spec = ElementSpec(

name, repeated, on_demand, identifier, namespace, attributes, children)

recursive = _str2bool(element_xml.get('recursive'))

if recursive:

element_spec.children[name] = element_spec

common_keys = set(element_spec.attributes).intersection(element_spec.children)

if common_keys:

raise RuntimeError(

'Element \'{}\' contains the following attributes and children with '

'the same name: \'{}\'. This violates the design assumptions of '

'this library. Please file a bug report. Thank you.'

.format(name, sorted(common_keys)))

return element_spec

def _parse_attribute(attribute_xml):

"""Parses an <attribute> element in the schema."""

name = attribute_xml.get('name')

required = _str2bool(attribute_xml.get('required'))

conflict_allowed = _str2bool(attribute_xml.get('conflict_allowed'))

conflict_behavior = attribute_xml.get('conflict_behavior', 'replace')

attribute_type = attribute_xml.get('type')

other_kwargs = {}

if attribute_type == 'keyword':

attribute_callable = attribute.Keyword

other_kwargs['valid_values'] = attribute_xml.get('valid_values').split(' ')

elif attribute_type == 'array':

array_size_str = attribute_xml.get('array_size')

attribute_callable = attribute.Array

other_kwargs['length'] = int(array_size_str) if array_size_str else None

other_kwargs['dtype'] = _ARRAY_DTYPE_MAP[attribute_xml.get('array_type')]

elif attribute_type == 'identifier':

attribute_callable = attribute.Identifier

elif attribute_type == 'reference':

attribute_callable = attribute.Reference

other_kwargs['reference_namespace'] = (

attribute_xml.get('reference_namespace') or name)

elif attribute_type == 'basepath':

attribute_callable = attribute.BasePath

other_kwargs['path_namespace'] = attribute_xml.get('path_namespace')

elif attribute_type == 'file':

attribute_callable = attribute.File

other_kwargs['path_namespace'] = attribute_xml.get('path_namespace')

else:

try:

attribute_callable = _SCALAR_TYPE_MAP[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,

conflict_allowed=conflict_allowed, conflict_behavior=conflict_behavior,

other_kwargs=other_kwargs)

def collect_namespaces(root_spec):

"""Constructs a set of namespaces in a given ElementSpec.

Args:

root_spec: An `ElementSpec` for the root element in the schema.

Returns:

A set of strings specifying the names of all the namespaces that are present

in the spec.

"""

findable_namespaces = set()

def update_namespaces_from_spec(spec):

findable_namespaces.add(spec.namespace)

for child_spec in spec.children.values():

if child_spec is not spec:

update_namespaces_from_spec(child_spec)

update_namespaces_from_spec(root_spec)

return findable_namespaces

MUJOCO = parse_schema(_SCHEMA_XML_PATH)

FINDABLE_NAMESPACES = frozenset(

collect_namespaces(MUJOCO).union(_ADDITIONAL_FINDABLE_NAMESPACES))

def _attachment_frame_spec(is_world_attachment):

"""Create specs for attachment frames.

Attachment frames are specialized <body> without an identifier.

The only allowed children are joints which also don't have identifiers.

Args:

is_world_attachment: Whether we are creating a spec for attachments to

worldbody. If `True`, allow <freejoint> as child.

Returns:

An `ElementSpec`.

"""

frame_spec = ElementSpec(

'body', repeated=True, on_demand=False, identifier=None, namespace='body',

attributes=collections.OrderedDict(),

children=collections.OrderedDict())

body_spec = MUJOCO.children['worldbody'].children['body']

# 'name' and 'childclass' attributes are excluded.

for attrib_name in (

'mocap', 'pos', 'quat', 'axisangle', 'xyaxes', 'zaxis', 'euler', 'user'):

frame_spec.attributes[attrib_name] = copy.deepcopy(

body_spec.attributes[attrib_name])

inertial_spec = body_spec.children['inertial']

frame_spec.children['inertial'] = copy.deepcopy(inertial_spec)

joint_spec = body_spec.children['joint']

frame_spec.children['joint'] = ElementSpec(

'joint', repeated=True, on_demand=False,

identifier=None, namespace='joint',

attributes=copy.deepcopy(joint_spec.attributes),

children=collections.OrderedDict())

if is_world_attachment:

freejoint_spec = (MUJOCO.children['worldbody']

.children['body'].children['freejoint'])

frame_spec.children['freejoint'] = ElementSpec(

'freejoint', repeated=False, on_demand=True,

identifier=None, namespace='joint',

attributes=copy.deepcopy(freejoint_spec.attributes),

children=collections.OrderedDict())

return frame_spec

ATTACHMENT_FRAME = _attachment_frame_spec(is_world_attachment=False)

WORLD_ATTACHMENT_FRAME = _attachment_frame_spec(is_world_attachment=True)

def override_schema(schema_xml_path):

"""Override the schema with a custom xml.

This method updates several global variables and care should be taken not to

call it if the pre-update values have already been used.

Args:

schema_xml_path: Path to schema xml file.

"""

global MUJOCO

global FINDABLE_NAMESPACES

MUJOCO = parse_schema(schema_xml_path)

FINDABLE_NAMESPACES = frozenset(

collect_namespaces(MUJOCO).union(_ADDITIONAL_FINDABLE_NAMESPACES))