PyPI - imt-ring - Versions diffs - 1.2.1__py3-none-any.whl - Mend

imt-ring 1.2.1__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.

Files changed (83) hide show

imt_ring-1.2.1.dist-info/METADATA +91 -0
imt_ring-1.2.1.dist-info/RECORD +83 -0
imt_ring-1.2.1.dist-info/WHEEL +5 -0
imt_ring-1.2.1.dist-info/top_level.txt +1 -0
ring/__init__.py +63 -0
ring/algebra.py +100 -0
ring/algorithms/__init__.py +45 -0
ring/algorithms/_random.py +403 -0
ring/algorithms/custom_joints/__init__.py +6 -0
ring/algorithms/custom_joints/rr_imp_joint.py +69 -0
ring/algorithms/custom_joints/rr_joint.py +33 -0
ring/algorithms/custom_joints/suntay.py +424 -0
ring/algorithms/dynamics.py +345 -0
ring/algorithms/generator/__init__.py +25 -0
ring/algorithms/generator/base.py +414 -0
ring/algorithms/generator/batch.py +282 -0
ring/algorithms/generator/motion_artifacts.py +222 -0
ring/algorithms/generator/pd_control.py +182 -0
ring/algorithms/generator/randomize.py +119 -0
ring/algorithms/generator/transforms.py +410 -0
ring/algorithms/generator/types.py +36 -0
ring/algorithms/jcalc.py +840 -0
ring/algorithms/kinematics.py +202 -0
ring/algorithms/sensors.py +582 -0
ring/base.py +1046 -0
ring/io/__init__.py +9 -0
ring/io/examples/branched.xml +24 -0
ring/io/examples/exclude/knee_trans_dof.xml +26 -0
ring/io/examples/exclude/standard_sys.xml +106 -0
ring/io/examples/exclude/standard_sys_rr_imp.xml +106 -0
ring/io/examples/inv_pendulum.xml +14 -0
ring/io/examples/knee_flexible_imus.xml +22 -0
ring/io/examples/spherical_stiff.xml +11 -0
ring/io/examples/symmetric.xml +12 -0
ring/io/examples/test_all_1.xml +39 -0
ring/io/examples/test_all_2.xml +39 -0
ring/io/examples/test_ang0_pos0.xml +9 -0
ring/io/examples/test_control.xml +16 -0
ring/io/examples/test_double_pendulum.xml +14 -0
ring/io/examples/test_free.xml +11 -0
ring/io/examples/test_kinematics.xml +23 -0
ring/io/examples/test_morph_system/four_seg_seg1.xml +26 -0
ring/io/examples/test_morph_system/four_seg_seg3.xml +26 -0
ring/io/examples/test_randomize_position.xml +26 -0
ring/io/examples/test_sensors.xml +13 -0
ring/io/examples/test_three_seg_seg2.xml +23 -0
ring/io/examples.py +42 -0
ring/io/test_examples.py +6 -0
ring/io/xml/__init__.py +6 -0
ring/io/xml/abstract.py +300 -0
ring/io/xml/from_xml.py +299 -0
ring/io/xml/test_from_xml.py +56 -0
ring/io/xml/test_to_xml.py +31 -0
ring/io/xml/to_xml.py +94 -0
ring/maths.py +397 -0
ring/ml/__init__.py +33 -0
ring/ml/base.py +292 -0
ring/ml/callbacks.py +434 -0
ring/ml/ml_utils.py +272 -0
ring/ml/optimizer.py +149 -0
ring/ml/params/0x13e3518065c21cd8.pickle +0 -0
ring/ml/ringnet.py +279 -0
ring/ml/train.py +318 -0
ring/ml/training_loop.py +131 -0
ring/rendering/__init__.py +2 -0
ring/rendering/base_render.py +271 -0
ring/rendering/mujoco_render.py +222 -0
ring/rendering/vispy_render.py +340 -0
ring/rendering/vispy_visuals.py +290 -0
ring/sim2real/__init__.py +7 -0
ring/sim2real/sim2real.py +288 -0
ring/spatial.py +126 -0
ring/sys_composer/__init__.py +5 -0
ring/sys_composer/delete_sys.py +114 -0
ring/sys_composer/inject_sys.py +110 -0
ring/sys_composer/morph_sys.py +361 -0
ring/utils/__init__.py +21 -0
ring/utils/batchsize.py +51 -0
ring/utils/colab.py +48 -0
ring/utils/hdf5.py +198 -0
ring/utils/normalizer.py +56 -0
ring/utils/path.py +44 -0
ring/utils/utils.py +161 -0

ring/io/examples/test_three_seg_seg2.xml ADDED Viewed

@@ -0,0 +1,23 @@
+<x_xy model="test_three_seg_seg2">
+    <options gravity="0 0 9.81" dt="0.01"/>
+    <defaults>
+        <geom edge_color="black" color="1 0.8 0.7 1"/>
+    </defaults>
+    <worldbody>
+        <body name="seg2" joint="free" damping="5 5 5 25 25 25">
+            <geom type="box" mass="1" pos="0.5 0 0" dim="1 0.25 0.2"/>
+            <body name="seg1" joint="ry" damping="3">
+                <geom type="box" mass="1" pos="-0.5 0 0" dim="1 0.25 0.2"/>
+                <body name="imu1" joint="frozen" pos="-0.5 0 0.125">
+                    <geom type="box" mass="0.1" dim="0.2 0.2 0.05" color="orange"/>
+                </body>
+            </body>
+            <body name="seg3" joint="rz" pos="1 0 0" damping="3">
+                <geom type="box" mass="1" pos="0.5 0 0" dim="1 0.25 0.2"/>
+                <body name="imu2" joint="frozen" pos="0.5 0 -0.125">
+                    <geom type="box" mass="0.1" dim="0.2 0.2 0.05" color="orange"/>
+                </body>
+            </body>
+        </body>
+    </worldbody>
+</x_xy>

ring/io/examples.py ADDED Viewed

@@ -0,0 +1,42 @@
+from pathlib import Path
+from typing import Iterator
+from ring import base
+from ring.utils import parse_path
+EXAMPLES_DIR = Path(__file__).parent.joinpath("examples")
+FOLDERS = ["", "test_morph_system"]
+EXCLUDE_FOLDERS = ["exclude"]
+def load_example(name: str):
+    "Load example from examples dir."
+    xml_path = parse_path(EXAMPLES_DIR, name, extension="xml")
+    return base.System.from_xml(xml_path)
+def list_examples() -> list[str]:
+    import os
+    def list_of_examples_in_folder(folder):
+        return [ex.split(".")[0] for ex in os.listdir(folder)]
+    examples = []
+    for folder in FOLDERS:
+        example_folder = list_of_examples_in_folder(EXAMPLES_DIR.joinpath(folder))
+        if len(folder) > 0:
+            example_folder = [folder + "/" + ex for ex in example_folder]
+        examples += example_folder
+    # exclude subfolders from examples
+    examples = list(set(examples) - set(FOLDERS) - set(EXCLUDE_FOLDERS))
+    examples.sort()
+    return examples
+def list_load_examples() -> Iterator[base.System]:
+    for example in list_examples():
+        yield load_example(example)

ring/io/test_examples.py ADDED Viewed

@@ -0,0 +1,6 @@
+import ring
+def test_examples():
+    for example in ring.io.list_examples():
+        ring.io.load_example(example)

ring/io/xml/__init__.py ADDED Viewed

@@ -0,0 +1,6 @@
+from .from_xml import load_comments_from_str
+from .from_xml import load_comments_from_xml
+from .from_xml import load_sys_from_str
+from .from_xml import load_sys_from_xml
+from .to_xml import save_sys_to_str
+from .to_xml import save_sys_to_xml

ring/io/xml/abstract.py ADDED Viewed

@@ -0,0 +1,300 @@
+from typing import Tuple, TypeVar
+import jax
+import jax.numpy as jnp
+import numpy as np
+from ring import base
+T = TypeVar("T")
+ATTR = dict
+default_quat = jnp.array([1.0, 0, 0, 0])
+default_pos = jnp.zeros((3,))
+default_damping = lambda qd_size, **_: jnp.zeros((qd_size,))
+default_armature = lambda qd_size, **_: jnp.zeros((qd_size,))
+default_stiffness = lambda qd_size, **_: jnp.zeros((qd_size,))
+def default_zeropoint(q_size, link_typ: str, **_):
+    zeropoint = jnp.zeros((q_size))
+    if link_typ in ["spherical", "free", "cor"]:
+        # zeropoint then is unit quaternion and not zeros
+        zeropoint = zeropoint.at[0].set(1.0)
+    return zeropoint
+default_fns = {
+    "damping": default_damping,
+    "armature": default_armature,
+    "spring_stiff": default_stiffness,
+    "spring_zero": default_zeropoint,
+}
+class AbsDampArmaStiffZero:
+    @staticmethod
+    def from_xml(attr: ATTR, q_size: int, qd_size: int, link_typ: str) -> list:
+        return [
+            jnp.atleast_1d(
+                attr.get(
+                    key,
+                    default_fns[key](q_size=q_size, qd_size=qd_size, link_typ=link_typ),
+                )
+            )
+            for key in ["damping", "armature", "spring_stiff", "spring_zero"]
+        ]
+    @staticmethod
+    def to_xml(
+        element: T,
+        damping: jax.Array,
+        armature: jax.Array,
+        stiffness: jax.Array,
+        zeropoint: jax.Array,
+        q_size: int,
+        qd_size: int,
+        link_typ: str,
+    ):
+        for key, arr in zip(
+            ["damping", "armature", "spring_stiff", "spring_zero"],
+            [damping, armature, stiffness, zeropoint],
+        ):
+            if not _arr_equal(
+                arr, default_fns[key](q_size=q_size, qd_size=qd_size, link_typ=link_typ)
+            ):
+                element.set(key, _to_str(arr))
+class AbsMaxCoordOMC:
+    @staticmethod
+    def from_xml(attr: ATTR) -> base.MaxCoordOMC:
+        pos = attr.get("pos", default_pos)
+        marker_number = int(attr.get("pos_marker"))
+        cs_name = attr.get("name")
+        return base.MaxCoordOMC(cs_name, marker_number, pos)
+    @staticmethod
+    def to_xml(element: T, max_coord_omc: base.MaxCoordOMC) -> None:
+        if not _arr_equal(max_coord_omc.pos_marker_constant_offset, default_pos):
+            element.set("pos", _to_str(max_coord_omc.pos_marker_constant_offset))
+        element.set("name", max_coord_omc.coordinate_system_name)
+        element.set("pos_marker", _to_str(max_coord_omc.pos_marker_number))
+class AbsTrans:
+    @staticmethod
+    def from_xml(attr: ATTR) -> base.Transform:
+        pos = attr.get("pos", default_pos)
+        rot = _get_rotation(attr)
+        return base.Transform(pos, rot)
+    @staticmethod
+    def to_xml(element: T, t: base.Transform) -> None:
+        if not _arr_equal(t.pos, default_pos):
+            element.set("pos", _to_str(t.pos))
+        if not _arr_equal(t.rot, default_quat):
+            element.set("quat", _to_str(t.rot))
+class AbsPosMinMax:
+    @staticmethod
+    def from_xml(attr: ATTR, pos: jax.Array) -> Tuple[jax.Array, jax.Array]:
+        pos_min = attr.get("pos_min", None)
+        pos_max = attr.get("pos_max", None)
+        assert (pos_min is None and pos_max is None) or (
+            pos_min is not None and pos_max is not None
+        ), (
+            f"In link {attr.get('name', 'None')} found only one of `pos_min` "
+            "and `pos_max`, but requires either both or none"
+        )
+        if pos_min is not None:
+            assert not _arr_equal(
+                pos_min, pos_max
+            ), f"In link {attr.get('name', 'None')} "
+            " both `pos_min` and `pos_max` are identical, use `pos` instead."
+        if pos_min is None:
+            pos_min = pos_max = pos
+        return pos_min, pos_max
+    @staticmethod
+    def to_xml(element: T, pos_min: jax.Array, pos_max: jax.Array):
+        if _arr_equal(pos_min, pos_max):
+            return
+        element.set("pos_min", _to_str(pos_min))
+        element.set("pos_max", _to_str(pos_max))
+def _from_xml_geom_attr_processing(geom_attr: ATTR):
+    "Common processing used by all geometries"
+    mass = geom_attr["mass"]
+    trafo = AbsTrans.from_xml(geom_attr)
+    # convert arrays to tuple[float], because of `struct.field(False)`
+    # Otherwise jitted functions with `sys` input will error on second execution, since
+    # it can't compare the two vispy_color arrays.
+    color = geom_attr.get("color", None)
+    if isinstance(color, (jax.Array, np.ndarray)):
+        color = tuple(color.tolist())
+    edge_color = geom_attr.get("edge_color", None)
+    if isinstance(edge_color, (jax.Array, np.ndarray)):
+        edge_color = tuple(edge_color.tolist())
+    return mass, trafo, color, edge_color
+def _to_xml_geom_processing(element: T, geom: base.Geometry) -> None:
+    "Common processing used by all geometries"
+    AbsTrans.to_xml(element, geom.transform)
+    element.set("mass", _to_str(geom.mass))
+    if geom.color is not None:
+        element.set("color", _to_str(geom.color))
+    if geom.edge_color is not None:
+        element.set("edge_color", _to_str(geom.edge_color))
+    element.set("type", geometry_to_xml_identifier[type(geom)])
+class AbsGeomBox:
+    xml_geom_type: str = "box"
+    geometry: base.Geometry = base.Box
+    @staticmethod
+    def from_xml(geom_attr: ATTR, link_idx: int) -> base.Box:
+        mass, trafo, color, edge_color = _from_xml_geom_attr_processing(geom_attr)
+        dims = [geom_attr["dim"][i] for i in range(3)]
+        assert all([dim > 0.0 for dim in dims]), "Negative box dimensions"
+        return base.Box(mass, trafo, link_idx, color, edge_color, *dims)
+    @staticmethod
+    def to_xml(element: T, geom: base.Box) -> None:
+        _to_xml_geom_processing(element, geom)
+        dim = np.array([geom.dim_x, geom.dim_y, geom.dim_z])
+        element.set("dim", _to_str(dim))
+class AbsGeomSphere:
+    xml_geom_type: str = "sphere"
+    geometry: base.Geometry = base.Sphere
+    @staticmethod
+    def from_xml(geom_attr: ATTR, link_idx: int) -> base.Sphere:
+        mass, trafo, color, edge_color = _from_xml_geom_attr_processing(geom_attr)
+        radius = geom_attr["dim"].item()
+        assert radius > 0.0, "Negative sphere radius"
+        return base.Sphere(mass, trafo, link_idx, color, edge_color, radius)
+    @staticmethod
+    def to_xml(element: T, geom: base.Sphere) -> None:
+        _to_xml_geom_processing(element, geom)
+        dim = np.array([geom.radius])
+        element.set("dim", _to_str(dim))
+class AbsGeomCylinder:
+    xml_geom_type: str = "cylinder"
+    geometry: base.Geometry = base.Cylinder
+    @staticmethod
+    def from_xml(geom_attr: ATTR, link_idx: int) -> base.Cylinder:
+        mass, trafo, color, edge_color = _from_xml_geom_attr_processing(geom_attr)
+        dims = [geom_attr["dim"][i] for i in range(2)]
+        assert all([dim > 0.0 for dim in dims]), "Negative cylinder dimensions"
+        return base.Cylinder(mass, trafo, link_idx, color, edge_color, *dims)
+    @staticmethod
+    def to_xml(element: T, geom: base.Cylinder) -> None:
+        _to_xml_geom_processing(element, geom)
+        dim = np.array([geom.radius, geom.length])
+        element.set("dim", _to_str(dim))
+class AbsGeomCapsule:
+    xml_geom_type: str = "capsule"
+    geometry: base.Geometry = base.Capsule
+    @staticmethod
+    def from_xml(geom_attr: ATTR, link_idx: int) -> base.Capsule:
+        mass, trafo, color, edge_color = _from_xml_geom_attr_processing(geom_attr)
+        dims = [geom_attr["dim"][i] for i in range(2)]
+        assert all([dim > 0.0 for dim in dims]), "Negative capsule dimensions"
+        return base.Capsule(mass, trafo, link_idx, color, edge_color, *dims)
+    @staticmethod
+    def to_xml(element: T, geom: base.Capsule) -> None:
+        _to_xml_geom_processing(element, geom)
+        dim = np.array([geom.radius, geom.length])
+        element.set("dim", _to_str(dim))
+class AbsGeomXYZ:
+    xml_geom_type: str = "xyz"
+    geometry: base.Geometry = base.XYZ
+    @staticmethod
+    def from_xml(geom_attr: ATTR, link_idx: int) -> base.XYZ:
+        if "dim" in geom_attr:
+            dim = geom_attr["dim"]
+        else:
+            dim = 1.0
+        assert dim > 0, "Negative xyz dimensions"
+        return base.XYZ.create(link_idx, dim)
+    @staticmethod
+    def to_xml(element: T, geom: base.XYZ):
+        element.set("type", geometry_to_xml_identifier[type(geom)])
+        if geom.size != 1.0:
+            element.set("dim", _to_str(geom.size))
+_ags = [
+    AbsGeomBox,
+    AbsGeomSphere,
+    AbsGeomCylinder,
+    AbsGeomCapsule,
+    AbsGeomXYZ,
+]
+geometry_to_xml_identifier = {ag.geometry: ag.xml_geom_type for ag in _ags}
+xml_identifier_to_abstract = {ag.xml_geom_type: ag for ag in _ags}
+geometry_to_abstract = {ag.geometry: ag for ag in _ags}
+def _arr_equal(a, b):
+    return np.all(np.array_equal(a, b))
+def _get_rotation(attr: ATTR):
+    rot = attr.get("quat", None)
+    if rot is not None:
+        assert "euler" not in attr, "Can't specify both `quat` and `euler` in xml"
+    elif "euler" in attr:
+        # we use zyx convention but angles are given
+        # in x, y, z in the xml file
+        # thus flip the order
+        euler_xyz = jnp.deg2rad(attr["euler"])
+        rot = base.maths.quat_euler(jnp.flip(euler_xyz), convention="zyx")
+    else:
+        rot = default_quat
+    return rot
+def _to_str(obj):
+    if isinstance(obj, list) or isinstance(obj, tuple):
+        if all([isinstance(ele, float) for ele in obj]):
+            obj = np.array(obj)
+    if isinstance(obj, (np.ndarray, jnp.ndarray)):
+        if obj.ndim == 0:
+            return str(obj)
+        return " ".join([str(x) for x in obj])
+    else:
+        return str(obj)

ring/io/xml/from_xml.py ADDED Viewed

@@ -0,0 +1,299 @@
+from xml.etree import ElementTree
+import jax
+import numpy as np
+from ring import base
+from ring.algorithms import jcalc
+from ring.utils import parse_path
+from . import abstract
+def _find_assert_unique(tree: ElementTree, *keys):
+    assert len(keys) > 0
+    value = tree.findall(keys[0])
+    if len(value) == 0:
+        return None
+    assert len(value) == 1
+    if len(keys) == 1:
+        return value[0]
+    else:
+        return _find_assert_unique(value[0], *keys[1:])
+def _build_defaults_attributes(tree):
+    tags = ["geom", "body"]
+    default_attrs = {}
+    for tag in tags:
+        defaults_subtree = _find_assert_unique(tree, "defaults", tag)
+        if defaults_subtree is None:
+            attrs = {}
+        else:
+            attrs = defaults_subtree.attrib
+        default_attrs[tag] = attrs
+    return default_attrs
+def _assert_all_tags_attrs_valid(xml_tree):
+    valid_attrs = {
+        "x_xy": ["model"],
+        "options": ["gravity", "dt"],
+        "defaults": ["geom", "body"],
+        "worldbody": [],
+        "body": [
+            "name",
+            "pos",
+            "pos_min",
+            "pos_max",
+            "quat",
+            "euler",
+            "joint",
+            "armature",
+            "damping",
+            "spring_stiff",
+            "spring_zero",
+        ],
+        "geom": ["type", "mass", "pos", "dim", "quat", "euler", "color", "edge_color"],
+        "omc": ["name", "pos_marker", "pos"],
+    }
+    for subtree in xml_tree.iter():
+        assert subtree.tag in list([key for key in valid_attrs])
+        for attr in subtree.attrib:
+            assert attr in valid_attrs[subtree.tag], f"attr {attr} not a valid attr"
+def _mix_in_defaults(worldbody, default_attrs):
+    for subtree in worldbody.iter():
+        if subtree.tag not in ["body", "geom"]:
+            continue
+        tag = subtree.tag
+        attr = subtree.attrib
+        for default_attr in default_attrs[tag]:
+            if default_attr not in attr:
+                attr.update({default_attr: default_attrs[tag][default_attr]})
+def _convert_attrs_to_arrays(xml_tree):
+    for subtree in xml_tree.iter():
+        for k, v in subtree.attrib.items():
+            try:
+                array = [float(num) for num in v.split(" ")]
+            except:  # noqa: E722
+                continue
+            subtree.attrib[k] = np.squeeze(np.array(array))
+def _extract_geoms_from_body_xml(body, current_link_idx):
+    link_geoms = []
+    for geom_subtree in body.findall("geom"):
+        attr = geom_subtree.attrib
+        geom = abstract.xml_identifier_to_abstract[attr["type"]].from_xml(
+            attr, current_link_idx
+        )
+        link_geoms.append(geom)
+    return link_geoms
+def _extract_omc_from_body_xml(body):
+    omc = body.findall("omc")
+    if len(omc) == 0:
+        return None
+    elif len(omc) == 1:
+        return abstract.AbsMaxCoordOMC.from_xml(omc[0].attrib)
+    else:
+        raise Exception(
+            f"Body `{body.attrib['name']}` has two or more `<omc ../>` fields."
+        )
+def _initial_setup(xml_tree):
+    _assert_all_tags_attrs_valid(xml_tree)
+    _convert_attrs_to_arrays(xml_tree)
+    default_attrs = _build_defaults_attributes(xml_tree)
+    worldbody = _find_assert_unique(xml_tree, "worldbody")
+    _mix_in_defaults(worldbody, default_attrs)
+    return worldbody
+DEFAULT_GRAVITY = np.array([0, 0, 9.81])
+DEFAULT_DT = 0.01
+def load_sys_from_str(xml_str: str, seed: int = 1) -> base.System:
+    """Load system from string input.
+    Args:
+        xml_str (str): XML Presentation of the system.
+    Returns:
+        base.System: Loaded system.
+    """
+    xml_tree = ElementTree.fromstring(xml_str)
+    worldbody = _initial_setup(xml_tree)
+    # check that <x_xy model="..."> syntax is correct
+    assert xml_tree.tag == "x_xy", (
+        "The root element in the xml of a x_xy model must be `x_xy`."
+        " Look up the examples under  x_xy/io/examples/*.xml to get started"
+    )
+    model_name = xml_tree.attrib.get("model", None)
+    # default options
+    options = {"gravity": DEFAULT_GRAVITY, "dt": DEFAULT_DT}
+    options_xml = _find_assert_unique(xml_tree, "options")
+    options.update({} if options_xml is None else options_xml.attrib)
+    # convert scalar array to float
+    # if this is uncommented, it leads to `ConcretizationTypeError`s
+    # options["dt"] = float(options["dt"])
+    links = {}
+    link_parents = {}
+    link_names = {}
+    link_types = {}
+    geoms = {}
+    armatures = {}
+    dampings = {}
+    spring_stiffnesses = {}
+    spring_zeropoints = {}
+    omc = {}
+    global_link_idx = -1
+    def process_body(body: ElementTree, parent: int):
+        nonlocal global_link_idx
+        global_link_idx += 1
+        current_link_idx = global_link_idx
+        current_link_typ = body.attrib["joint"]
+        if current_link_typ == "cor":
+            raise Exception(
+                "`cor` joint type is not meant to be used like this. Either use a "
+                "`free` joint instead of `cor` and set MotionConfig.cor=True or, use"
+                " a free joint and call sys._replace_free_with_cor."
+            )
+        link_parents[current_link_idx] = parent
+        link_types[current_link_idx] = current_link_typ
+        link_names[current_link_idx] = body.attrib["name"]
+        transform = abstract.AbsTrans.from_xml(body.attrib)
+        pos_min, pos_max = abstract.AbsPosMinMax.from_xml(body.attrib, transform.pos)
+        links[current_link_idx] = base.Link(transform, pos_min, pos_max)
+        omc[current_link_idx] = _extract_omc_from_body_xml(body)
+        q_size = base.Q_WIDTHS[current_link_typ]
+        qd_size = base.QD_WIDTHS[current_link_typ]
+        (
+            damping,
+            armature,
+            stiffness,
+            zeropoint,
+        ) = abstract.AbsDampArmaStiffZero.from_xml(
+            body.attrib, q_size, qd_size, current_link_typ
+        )
+        armatures[current_link_idx] = armature
+        dampings[current_link_idx] = damping
+        spring_stiffnesses[current_link_idx] = stiffness
+        spring_zeropoints[current_link_idx] = zeropoint
+        geoms[current_link_idx] = _extract_geoms_from_body_xml(body, current_link_idx)
+        for subbodies in body.findall("body"):
+            process_body(subbodies, current_link_idx)
+        return
+    for body in worldbody.findall("body"):
+        process_body(body, -1)
+    def assert_order_then_to_list(d: dict) -> list:
+        assert [i for i in d] == list(range(len(d)))
+        return [d[i] for i in d]
+    links = assert_order_then_to_list(links)
+    links = links[0].batch(*links[1:])
+    dampings = np.concatenate(assert_order_then_to_list(dampings))
+    armatures = np.concatenate(assert_order_then_to_list(armatures))
+    spring_stiffnesses = np.concatenate(assert_order_then_to_list(spring_stiffnesses))
+    spring_zeropoints = np.concatenate(assert_order_then_to_list(spring_zeropoints))
+    # add all geoms directly connected to worldbody
+    flat_geoms = [geom for geoms in assert_order_then_to_list(geoms) for geom in geoms]
+    flat_geoms += _extract_geoms_from_body_xml(worldbody, -1)
+    sys = base.System(
+        link_parents=assert_order_then_to_list(link_parents),
+        links=links,
+        link_types=assert_order_then_to_list(link_types),
+        link_damping=dampings,
+        link_armature=armatures,
+        link_spring_stiffness=spring_stiffnesses,
+        link_spring_zeropoint=spring_zeropoints,
+        dt=float(options["dt"]),
+        geoms=flat_geoms,
+        gravity=options["gravity"],
+        link_names=assert_order_then_to_list(link_names),
+        model_name=model_name,
+        omc=assert_order_then_to_list(omc),
+    )
+    # numpy -> jax
+    # we load using numpy in order to have float64 precision
+    sys = jax.tree_map(jax.numpy.asarray, sys)
+    sys = jcalc._init_joint_params(jax.random.PRNGKey(seed), sys)
+    return sys.parse()
+def load_sys_from_xml(xml_path: str, seed: int = 1):
+    return load_sys_from_str(_load_xml(xml_path), seed=seed)
+def _load_xml(xml_path: str) -> str:
+    xml_path = parse_path(xml_path, extension="xml")
+    with open(xml_path, "r") as f:
+        xml_str = f.read()
+    return xml_str
+def load_comments_from_xml(xml_path: str, key: str) -> list[dict]:
+    """Example:
+    test.xml
+    <!--keyname1 key1=val1 key2=2-->
+    <!--keyname1 key1=val1 key2=3-->
+    <!--keyname2 key1=val1 key2=2-->
+    `load_comments_from_xml(test.xml, key=keyname1)`
+    Returns:
+    >>> [{key1: val1, key2: 2}, {key1: val1, key2: 3}]
+    """
+    return load_comments_from_str(_load_xml(xml_path), key=key)
+def load_comments_from_str(xml_str: str, key: str) -> list[dict]:
+    parser = ElementTree.XMLParser(target=ElementTree.TreeBuilder(insert_comments=True))
+    tree = ElementTree.fromstring(xml_str, parser)
+    comments = []
+    for node in tree.iter():
+        if "function Comment" in str(node.tag):
+            comments.append(node.text)
+    filtered = [s.split(" ")[1:] for s in comments if s.split(" ")[0] == key]
+    comments_dict = []
+    for comment in filtered:
+        d = dict()
+        for pair in comment:
+            key, value = pair.split("=")
+            d[key] = value
+        comments_dict.append(d)
+    return comments_dict