imt-ring 1.2.1__py3-none-any.whl

ring/sim2real/sim2real.py

@@ -0,0 +1,288 @@
+from typing import Optional, Tuple
+
+import jax
+from ring import algebra
+from ring import base
+from ring import io
+from ring import maths
+from ring.algorithms import generator
+from ring.algorithms import jcalc
+import tree_utils
+
+
+def xs_from_raw(
+    sys: base.System,
+    link_name_pos_rot: dict,
+    eps_frame: Optional[str] = None,
+    qinv: bool = False,
+) -> base.Transform:
+    """Build time-series of maximal coordinates `xs` from raw position and
+    quaternion trajectory data. This function scans through each link (as
+    defined by `sys`), looks for the raw data in `link_name_pos_rot` using
+    the `link_name` as identifier. It inverts the quaternion if `qinv`.
+    Then, it creates a `Transform` that transforms from epsilon (as defined
+    by `eps_frame`) to the link for each timestep. Finally, it stacks all
+    transforms in order as defined by `sys` along the 1-th axis. The 0-th
+    axis is time axis.
+
+    Args:
+        sys (ring.base.System): System which defines ordering of returned `xs`
+        link_name_pos_rot (dict): Dictonary of `link_name` ->
+            {'pos': ..., 'quat': ...}. Obtained, e.g., using `process_omc`.
+        eps_frame (str, optional): Move into this segment's frame at time zero as
+            eps frame. Defaults to `None`.
+            If `None`: Don't move into a specific eps-frame.
+
+    Returns:
+        ring.base.Transform: Time-series of eps-to-link transformations
+    """
+    # determine `eps_frame` transform
+    if eps_frame is not None:
+        eps = link_name_pos_rot[eps_frame]
+        q_eps = eps["quat"][0]
+        if qinv:
+            q_eps = maths.quat_inv(q_eps)
+        t_eps = base.Transform(eps["pos"][0], q_eps)
+    else:
+        t_eps = base.Transform.zero()
+
+    # build `xs` from optical motion capture data
+    xs = []
+
+    def f(_, __, link_name: str):
+        q = link_name_pos_rot[link_name]["quat"]
+        pos = link_name_pos_rot[link_name].get("pos", None)
+        if qinv:
+            q = maths.quat_inv(q)
+        t = base.Transform.create(pos, q)
+        t = algebra.transform_mul(t, algebra.transform_inv(t_eps))
+        xs.append(t)
+
+    sys.scan(f, "l", sys.link_names)
+
+    # stack and permute such that time-axis is 0-th axis
+    xs = xs[0].batch(*xs[1:])
+    xs = xs.transpose((1, 0, 2))
+    return xs
+
+
+def match_xs(
+    sys: base.System, xs: base.Transform, sys_xs: base.System
+) -> base.Transform:
+    """Match tranforms `xs` to subsystem `sys`.
+
+    Args:
+        sys (System): Smaller system. Every link in `sys` must be in `sys_xs`.
+        xs (Transform): Transforms of larger system.
+        sys_xs (Transform): Larger system.
+
+    Returns:
+        Transform: Transforms of smaller system.
+    """
+    _checks_time_series_of_xs(sys_xs, xs)
+
+    xs_small = xs_from_raw(
+        sys,
+        {
+            name: {
+                "pos": xs.pos[:, sys_xs.name_to_idx(name)],
+                "quat": xs.rot[:, sys_xs.name_to_idx(name)],
+            }
+            for name in sys_xs.link_names
+        },
+        eps_frame=None,
+        qinv=False,
+    )
+    return xs_small
+
+
+def unzip_xs(
+    sys: base.System, xs: base.Transform
+) -> Tuple[base.Transform, base.Transform]:
+    """Split eps-to-link transforms into parent-to-child pure
+    translational `transform1` and pure rotational `transform2`.
+
+    Args:
+        sys (System): Defines scan.tree
+        xs (Transform): Eps-to-link transforms
+
+    Returns:
+        Tuple[Transform, Transform]: transform1, transform2
+    """
+    _checks_time_series_of_xs(sys, xs)
+
+    @jax.vmap
+    def _unzip_xs(xs):
+        def f(_, __, i: int, p: int):
+            if p == -1:
+                x_parent_to_link = xs[i]
+            else:
+                x_parent_to_link = algebra.transform_mul(
+                    xs[i], algebra.transform_inv(xs[p])
+                )
+
+            transform1_pos = base.Transform.create(pos=x_parent_to_link.pos)
+            transform2_rot = base.Transform.create(rot=x_parent_to_link.rot)
+            return (transform1_pos, transform2_rot)
+
+        return sys.scan(f, "ll", list(range(sys.num_links())), sys.link_parents)
+
+    return _unzip_xs(xs)
+
+
+def zip_xs(
+    sys: base.System,
+    xs_transform1: base.Transform,
+    xs_transform2: base.Transform,
+) -> base.Transform:
+    """Performs forward kinematics using `transform1` and `transform2`.
+
+    Args:
+        sys (ring.base.System): Defines scan_sys
+        xs_transform1 (ring.base.Transform): Applied before `transform1`
+        xs_transform2 (ring.base.Transform): Applied after `transform2`
+
+    Returns:
+        ring.base.Transform: Time-series of eps-to-link transformations
+    """
+    _checks_time_series_of_xs(sys, xs_transform1)
+    _checks_time_series_of_xs(sys, xs_transform2)
+
+    @jax.vmap
+    def _zip_xs(xs_transform1, xs_transform2):
+        eps_to_l = {-1: base.Transform.zero()}
+
+        def f(_, __, i: int, p: int):
+            transform = algebra.transform_mul(xs_transform2[i], xs_transform1[i])
+            eps_to_l[i] = algebra.transform_mul(transform, eps_to_l[p])
+            return eps_to_l[i]
+
+        return sys.scan(f, "ll", list(range(sys.num_links())), sys.link_parents)
+
+    return _zip_xs(xs_transform1, xs_transform2)
+
+
+def _checks_time_series_of_xs(sys, xs):
+    assert tree_utils.tree_ndim(xs) == 3, f"pos.shape={xs.pos.shape}"
+    num_links_xs, num_links_sys = tree_utils.tree_shape(xs, axis=1), sys.num_links()
+    assert num_links_xs == num_links_sys, f"{num_links_xs} != {num_links_sys}"
+
+
+def delete_to_world_pos_rot(sys: base.System, xs: base.Transform) -> base.Transform:
+    """Replace the transforms of all links that connect to the worldbody
+    by unity transforms.
+
+    Args:
+        sys (System): System only used for structure (in scan_sys).
+        xs (Transform): Time-series of transforms to be modified.
+
+    Returns:
+        Transform: Time-series of modified transforms.
+    """
+    _checks_time_series_of_xs(sys, xs)
+
+    zero_trafo = base.Transform.zero((xs.shape(),))
+    for i, p in enumerate(sys.link_parents):
+        if p == -1:
+            xs = _overwrite_transform_of_link_then_update(sys, xs, zero_trafo, i)
+    return xs
+
+
+def randomize_to_world_pos_rot(
+    key: jax.Array, sys: base.System, xs: base.Transform, config: jcalc.MotionConfig
+) -> base.Transform:
+    """Replace the transforms of all links that connect to the worldbody
+    by randomize transforms.
+
+    Args:
+        key (jax.Array): PRNG Key.
+        sys (System): System only used for structure (in scan_sys).
+        xs (Transform): Time-series of transforms to be modified.
+        config (MotionConfig): Defines the randomization.
+
+    Returns:
+        Transform: Time-series of modified transforms.
+    """
+    _checks_time_series_of_xs(sys, xs)
+    assert sys.link_parents.count(-1) == 1, "Found multiple connections to world"
+
+    free_sys_str = """
+<x_xy>
+    <options dt="0.01"/>
+    <worldbody>
+        <body name="free" joint="free"/>
+    </worldbody>
+</x_xy>
+"""
+
+    free_sys = io.load_sys_from_str(free_sys_str)
+    _, xs_free = generator.RCMG(
+        free_sys, config, finalize_fn=lambda key, q, x, sys: (q, x)
+    ).to_lazy_gen()(key)
+    xs_free = xs_free.take(0, axis=0)
+    xs_free = xs_free.take(free_sys.name_to_idx("free"), axis=1)
+    link_idx_to_world = sys.link_parents.index(-1)
+    return _overwrite_transform_of_link_then_update(sys, xs, xs_free, link_idx_to_world)
+
+
+def _overwrite_transform_of_link_then_update(
+    sys: base.System, xs: base.Transform, xs_new_link: base.Transform, new_link_idx: int
+):
+    """Replace transform and then perform forward kinematics."""
+    assert xs_new_link.ndim() == (xs.ndim() - 1) == 2
+    transform1, transform2 = unzip_xs(sys, xs)
+    transform1 = _replace_transform_of_link(transform1, xs_new_link, new_link_idx)
+    zero_trafo = base.Transform.zero((xs_new_link.shape(),))
+    transform2 = _replace_transform_of_link(transform2, zero_trafo, new_link_idx)
+    return zip_xs(sys, transform1, transform2)
+
+
+def _replace_transform_of_link(
+    xs: base.Transform, xs_new_link: base.Transform, link_idx
+):
+    return xs.transpose((1, 0, 2)).index_set(link_idx, xs_new_link).transpose((1, 0, 2))
+
+
+def scale_xs(
+    sys: base.System,
+    xs: base.Transform,
+    factor: float,
+    exclude: list[str] = ["px", "py", "pz", "free"],
+) -> base.Transform:
+    """Increase / decrease transforms by scaling their positional / rotational
+    components based on the systems link type, i.e. the `xs` should conceptionally
+    be `transform2` objects.
+
+    Args:
+        sys (System): System defining structure (for scan_sys)
+        xs (Transform): Time-series of transforms to be modified.
+        factor (float): Multiplicative factor.
+        exclude (list[str], optional): Skip scaling of transforms if their link_type
+            is one of those. Defaults to ["px", "py", "pz", "free"].
+
+    Returns:
+        Transform: Time-series of scaled transforms.
+    """
+    _checks_time_series_of_xs(sys, xs)
+
+    @jax.vmap
+    def _scale_xs(xs):
+        def f(_, __, i: int, type: str):
+            x_link = xs[i]
+            if type not in exclude:
+                x_link = _scale_transform_based_on_type(x_link, type, factor)
+            return x_link
+
+        return sys.scan(f, "ll", list(range(sys.num_links())), sys.link_types)
+
+    return _scale_xs(xs)
+
+
+def _scale_transform_based_on_type(x: base.Transform, link_type: str, factor: float):
+    pos, rot = x.pos, x.rot
+    if link_type in ["px", "py", "pz", "free"]:
+        pos = pos * factor
+    if link_type in ["rx", "ry", "rz", "spherical", "free"]:
+        axis, angle = maths.quat_to_rot_axis(rot)
+        rot = maths.quat_rot_axis(axis, angle * factor)
+    return base.Transform(pos, rot)

ring/spatial.py

@@ -0,0 +1,126 @@
+"""
+Implements `Table 1` from
+    `A Beginner's Guide to 6-D Vectors (Part 2)`
+    by Roy Featherstone.
+
+`A small but sufficient set of spatial arithemtic operations.`
+"""
+
+
+import jax.numpy as jnp
+
+
+def rx(theta):
+    """
+    [
+        1  0 0
+        0  c s
+        0 -s c
+    ]
+    where c = cos(theta)
+          s = sin(theta)
+    """
+    s, c = jnp.sin(theta), jnp.cos(theta)
+    return jnp.array([[1, 0, 0], [0, c, s], [0, -s, c]])
+
+
+def ry(theta):
+    """
+    [
+        c 0 -s
+        0 1  0
+        s 0  c
+    ]
+    where c = cos(theta)
+          s = sin(theta)
+    """
+    s, c = jnp.sin(theta), jnp.cos(theta)
+    return jnp.array([[c, 0, -s], [0, 1, 0], [s, 0, c]])
+
+
+def rz(theta):
+    """
+    [
+         c s 0
+        -s c 0
+         0 0 1
+    ]
+    where c = cos(theta)
+          s = sin(theta)
+    """
+    s, c = jnp.sin(theta), jnp.cos(theta)
+    return jnp.array([[c, s, 0], [-s, c, 0], [0, 0, 1]])
+
+
+def cross(r):
+    assert r.shape == (3,)
+    return jnp.array([[0, -r[2], r[1]], [r[2], 0, -r[0]], [-r[1], r[0], 0]])
+
+
+def quadrants(aa=None, ab=None, ba=None, bb=None, default=jnp.zeros):
+    M = default((6, 6))
+    if aa is not None:
+        M = M.at[:3, :3].set(aa)
+    if ab is not None:
+        M = M.at[:3, 3:].set(ab)
+    if ba is not None:
+        M = M.at[3:, :3].set(ba)
+    if bb is not None:
+        M = M.at[3:, 3:].set(bb)
+    return M
+
+
+def crm(v):
+    assert v.shape == (6,)
+    return quadrants(cross(v[:3]), ba=cross(v[3:]), bb=cross(v[:3]))
+
+
+def crf(v):
+    return -crm(v).T
+
+
+def _rotxyz(E):
+    return quadrants(E, bb=E)
+
+
+def rotx(theta):
+    return _rotxyz(rx(theta))
+
+
+def roty(theta):
+    return _rotxyz(ry(theta))
+
+
+def rotz(theta):
+    return _rotxyz(rz(theta))
+
+
+def xlt(r):
+    assert r.shape == (3,)
+    return quadrants(jnp.eye(3), ba=-cross(r), bb=jnp.eye(3))
+
+
+def X_transform(E, r):
+    return _rotxyz(E) @ xlt(r)
+
+
+def mcI(m, c, Ic):
+    assert c.shape == (3,)
+    assert Ic.shape == (3, 3)
+    return quadrants(
+        Ic - m * cross(c) @ cross(c), m * cross(c), -m * cross(c), m * jnp.eye(3)
+    )
+
+
+def XtoV(X):
+    assert X.shape == (6, 6)
+    return 0.5 * jnp.array(
+        [
+            [X[1, 2] - X[2, 1]],
+            [X[2, 0] - X[0, 2]],
+            [X[0, 1] - X[1, 0]],
+            [X[4, 2] - X[5, 1]],
+            [X[5, 0] - X[3, 2]],
+            [X[3, 1] - X[4, 0]],
+        ]
+    )

ring/sys_composer/__init__.py

@@ -0,0 +1,5 @@
+from .delete_sys import delete_subsystem
+from .delete_sys import make_sys_noimu
+from .inject_sys import inject_system
+from .morph_sys import identify_system
+from .morph_sys import morph_system

ring/sys_composer/delete_sys.py

@@ -0,0 +1,114 @@
+from typing import Optional
+
+import jax.numpy as jnp
+from ring import base
+import tree_utils
+
+
+def _autodetermine_imu_names(sys: base.System) -> list[str]:
+    return sys.findall_imus()
+
+
+def make_sys_noimu(sys: base.System, imu_link_names: Optional[list[str]] = None):
+    "Returns, e.g., imu_attachment = {'imu1': 'seg1', 'imu2': 'seg3'}"
+    if imu_link_names is None:
+        imu_link_names = _autodetermine_imu_names(sys)
+    imu_attachment = {name: sys.parent_name(name) for name in imu_link_names}
+    sys_noimu = delete_subsystem(sys, imu_link_names)
+    return sys_noimu, imu_attachment
+
+
+def delete_subsystem(
+    sys: base.System, link_name: str | list[str], strict: bool = True
+) -> base.System:
+    "Cut subsystem starting at `link_name` (inclusive) from tree."
+    if isinstance(link_name, list):
+        for ln in link_name:
+            sys = delete_subsystem(sys, ln, strict)
+        return sys
+
+    if not strict:
+        try:
+            return delete_subsystem(sys, link_name, strict=True)
+        except AssertionError:
+            return sys
+
+    assert (
+        link_name in sys.link_names
+    ), f"link {link_name} not found in {sys.link_names}"
+
+    subsys = _find_subsystem_indices(sys.link_parents, sys.name_to_idx(link_name))
+    idx_map, keep = _idx_map_and_keepers(sys.link_parents, subsys)
+
+    def take(list):
+        return [ele for i, ele in enumerate(list) if i in keep]
+
+    d, a, ss, sz = [], [], [], []
+
+    def filter_arrays(_, __, damp, arma, stiff, zero, i: int):
+        if i in keep:
+            d.append(damp)
+            a.append(arma)
+            ss.append(stiff)
+            sz.append(zero)
+
+    sys.scan(
+        filter_arrays,
+        "dddql",
+        sys.link_damping,
+        sys.link_armature,
+        sys.link_spring_stiffness,
+        sys.link_spring_zeropoint,
+        list(range(sys.num_links())),
+    )
+
+    d, a, ss, sz = map(jnp.concatenate, (d, a, ss, sz))
+
+    new_sys = base.System(
+        link_parents=_reindex_parent_array(sys.link_parents, subsys),
+        links=tree_utils.tree_indices(sys.links, jnp.array(keep, dtype=int)),
+        link_types=take(sys.link_types),
+        link_damping=d,
+        link_armature=a,
+        link_spring_stiffness=ss,
+        link_spring_zeropoint=sz,
+        dt=sys.dt,
+        geoms=[
+            geom.replace(link_idx=idx_map[geom.link_idx])
+            for geom in sys.geoms
+            if geom.link_idx in keep
+        ],
+        gravity=sys.gravity,
+        integration_method=sys.integration_method,
+        mass_mat_iters=sys.mass_mat_iters,
+        link_names=take(sys.link_names),
+        model_name=sys.model_name,
+        omc=take(sys.omc),
+    )
+
+    return new_sys.parse()
+
+
+def _find_subsystem_indices(parents: list[int], k: int) -> list[int]:
+    subsys = [k]
+    for i, p in enumerate(parents):
+        if p in subsys:
+            subsys.append(i)
+    return subsys
+
+
+def _idx_map_and_keepers(parents: list[int], subsys: list[int]):
+    num_links = len(parents)
+    # keep must be in ascending order
+    keep = []
+    for i in range(num_links):
+        if i not in subsys:
+            keep.append(i)
+
+    idx_map = dict(zip([-1] + keep, range(-1, len(keep))))
+    return idx_map, keep
+
+
+def _reindex_parent_array(parents: list[int], subsys: list[int]) -> list[int]:
+    idx_map, keep = _idx_map_and_keepers(parents, subsys)
+    return [idx_map[p] for i, p in enumerate(parents) if i in keep]

ring/sys_composer/inject_sys.py

@@ -0,0 +1,110 @@
+from typing import Optional
+
+import jax
+import jax.numpy as jnp
+from ring import base
+from tree_utils import tree_batch
+
+
+def _tree_nan_like(tree, repeats: int):
+    return jax.tree_map(
+        lambda arr: jnp.repeat(arr[0:1] * jnp.nan, repeats, axis=0), tree
+    )
+
+
+# TODO
+# right now this function won't really keep index ordering
+# as one might expect.
+# It simply appends the `sub_sys` at index end of `sys`, even
+# though it might be injected in the middle of the index range
+# of `sys`.
+# This will be fixed once we have a `dump_sys_to_xml` function
+def inject_system(
+    sys: base.System,
+    sub_sys: base.System,
+    at_body: Optional[str] = None,
+) -> base.System:
+    """Combine two systems into one.
+
+    Args:
+        sys (base.System): Large system.
+        sub_sys (base.System): Small system that will be included into the
+            large system `sys`.
+        at_body (Optional[str], optional): Into which body of the large system
+            small system will be included. Defaults to `worldbody`.
+
+    Returns:
+        base.System: _description_
+    """
+
+    # replace parent array
+    if at_body is None:
+        new_world = -1
+    else:
+        new_world = sys.name_to_idx(at_body)
+
+    # append sub_sys at index end and replace sub_sys world with `at_body`
+    N = sys.num_links()
+
+    def new_parent(old_parent: int):
+        if old_parent != -1:
+            return old_parent + N
+        else:
+            return new_world
+
+    sub_sys = sub_sys.replace(
+        link_parents=[new_parent(p) for p in sub_sys.link_parents]
+    )
+
+    # replace link indices of geoms in sub_sys
+    sub_sys = sub_sys.replace(
+        geoms=[
+            geom.replace(link_idx=new_parent(geom.link_idx)) for geom in sub_sys.geoms
+        ]
+    )
+
+    # build union of two joint_params dictionaries because each system might have custom
+    # joints that the other does not have
+    missing_in_sys = set(sub_sys.links.joint_params.keys()) - set(
+        sys.links.joint_params.keys()
+    )
+    sys_n_links = sys.num_links()
+    for typ in missing_in_sys:
+        sys.links.joint_params[typ] = _tree_nan_like(
+            sub_sys.links.joint_params[typ], sys_n_links
+        )
+
+    missing_in_subsys = set(
+        sys.links.joint_params.keys() - sub_sys.links.joint_params.keys()
+    )
+    subsys_n_links = sub_sys.num_links()
+    for typ in missing_in_subsys:
+        sub_sys.links.joint_params[typ] = _tree_nan_like(
+            sys.links.joint_params[typ], subsys_n_links
+        )
+
+    # merge two systems
+    concat = lambda a1, a2: tree_batch([a1, a2], True, "jax")
+    combined_sys = base.System(
+        link_parents=sys.link_parents + sub_sys.link_parents,
+        links=concat(sys.links, sub_sys.links),
+        link_types=sys.link_types + sub_sys.link_types,
+        link_damping=concat(sys.link_damping, sub_sys.link_damping),
+        link_armature=concat(sys.link_armature, sub_sys.link_armature),
+        link_spring_stiffness=concat(
+            sys.link_spring_stiffness, sub_sys.link_spring_stiffness
+        ),
+        link_spring_zeropoint=concat(
+            sys.link_spring_zeropoint, sub_sys.link_spring_zeropoint
+        ),
+        dt=sys.dt,
+        geoms=sys.geoms + sub_sys.geoms,
+        gravity=sys.gravity,
+        integration_method=sys.integration_method,
+        mass_mat_iters=sys.mass_mat_iters,
+        link_names=sys.link_names + sub_sys.link_names,
+        model_name=sys.model_name,
+        omc=sys.omc + sub_sys.omc,
+    )
+
+    return combined_sys.parse()