Source code for rofunc.utils.robolab.formatter.mjcf_parser.export_with_assets

# 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.
# ============================================================================

"""Saves Mujoco models with relevant assets."""

import os

from rofunc.utils.robolab.formatter.mjcf_parser import util

from rofunc.utils.robolab.formatter.mjcf_parser import constants


def export_with_assets(mjcf_model, out_dir, out_file_name=None,
                       *,
                       precision=constants.XML_DEFAULT_PRECISION,
                       zero_threshold=0):
    """Saves mjcf.model in the given directory in MJCF (XML) format.

    Creates an MJCF XML file named `out_file_name` in the specified `out_dir`, and
    writes all of its assets into the same directory.

    Args:
      mjcf_model: `mjcf.RootElement` instance to export.
      out_dir: Directory to save the model and assets. Will be created if it does
        not already exist.
      out_file_name: (Optional) Name of the XML file to create. Defaults to the
        model name (`mjcf_model.model`) suffixed with '.xml'.
      precision: (optional) Number of digits to output for floating point
        quantities.
      zero_threshold: (optional) When outputting XML, floating point quantities
        whose absolute value falls below this threshold will be treated as zero.

    Raises:
      ValueError: If `out_file_name` is a string that does not end with '.xml'.
    """
    if out_file_name is None:
        out_file_name = mjcf_model.model + '.xml'
    elif not out_file_name.lower().endswith('.xml'):
        raise ValueError('If `out_file_name` is specified it must end with '
                         '\'.xml\': got {}'.format(out_file_name))
    assets = mjcf_model.get_assets()
    # This should never happen because `mjcf` does not support `.xml` assets.
    assert out_file_name not in assets
    assets[out_file_name] = mjcf_model.to_xml_string(
        precision=precision, zero_threshold=zero_threshold)
    if not os.path.exists(out_dir):
        os.makedirs(out_dir)
    for filename, contents in assets.items():
        with open(os.path.join(out_dir, filename), 'wb') as f:
            f.write(util.to_binary_string(contents))