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/algorithms/generator/motion_artifacts.py ADDED Viewed

@@ -0,0 +1,222 @@
+import warnings
+import jax
+import jax.numpy as jnp
+from ring import base
+from ring import io
+import tree_utils
+def imu_reference_link_name(imu_link_name: str) -> str:
+    return "_" + imu_link_name
+def unactuated_subsystem(sys) -> list[str]:
+    return [imu_reference_link_name(name) for name in sys.findall_imus()]
+def _subsystem_factory(imu_name: str, pos_min_max: float) -> base.System:
+    assert pos_min_max >= 0
+    pos = f'pos_min="-{pos_min_max} -{pos_min_max} -{pos_min_max}" pos_max="{pos_min_max} {pos_min_max} {pos_min_max}"'  # noqa: E501
+    stiff = 'spring_stiff="50 50 50"'
+    damping = 'damping="5 5 5"'
+    return io.load_sys_from_str(
+        f"""
+        <x_xy>
+        <worldbody>
+        <body name="{imu_name}" joint="p3d" {pos if pos_min_max != 0.0 else ""} {stiff} {damping}/>
+        </worldbody>
+        </x_xy>
+        """  # noqa: E501
+    )
+def inject_subsystems(
+    sys: base.System,
+    pos_min_max: float = 0.0,
+    **kwargs,
+) -> base.System:
+    imu_idx_to_name_map = {sys.name_to_idx(imu): imu for imu in sys.findall_imus()}
+    default_spher_stif = jnp.ones((3,)) * 0.3
+    default_spher_damp = default_spher_stif * 0.1
+    for imu in sys.findall_imus():
+        sys = sys.unfreeze(imu, "spherical")
+        # set default stiffness and damping of spherical joint
+        # this won't override anything because the frozen joint can not have any values
+        qd_slice = sys.idx_map("d")[imu]
+        stiffne = sys.link_spring_stiffness.at[qd_slice].set(default_spher_stif)
+        damping = sys.link_damping.at[qd_slice].set(default_spher_damp)
+        sys = sys.replace(link_spring_stiffness=stiffne, link_damping=damping)
+        _imu = imu_reference_link_name(imu)
+        sys = sys.change_link_name(imu, _imu)
+        sys = sys.inject_system(_subsystem_factory(imu, pos_min_max), _imu)
+    # attach geoms to newly injected link
+    new_geoms = []
+    for geom in sys.geoms:
+        if geom.link_idx in imu_idx_to_name_map:
+            imu_name = imu_idx_to_name_map[geom.link_idx]
+            new_link_idx = sys.name_to_idx(imu_name)
+            geom = geom.replace(link_idx=new_link_idx)
+        new_geoms.append(geom)
+    sys = sys.replace(geoms=new_geoms)
+    # TODO investigate whether this parse is needed; I don't think so
+    # re-calculate the inertia matrices because the geoms have been re-attached
+    sys = sys.parse()
+    # TODO set all joint_params to zeros; they can not be preserved anyways and
+    # otherwise many warnings will be rose
+    # instead warn explicitly once now and move on
+    warnings.warn(
+        "`sys.links.joint_params` has been set to zero, this might lead to "
+        "unexpected behaviour unless you use `randomize_joint_params`"
+    )
+    joint_params_zeros = tree_utils.tree_zeros_like(sys.links.joint_params)
+    sys = sys.replace(links=sys.links.replace(joint_params=joint_params_zeros))
+    # double load; this fixes the issue that injected links got appended at the end
+    sys = io.load_sys_from_str(io.save_sys_to_str(sys))
+    return sys
+_STIF_MIN_SPH = 0.2
+_STIF_MAX_SPH = 10.0
+_STIF_MIN_P3D = 25.0
+_STIF_MAX_P3D = 1e3
+# damping = factor * stiffness
+_DAMP_MIN = 0.05
+_DAMP_MAX = 0.5
+def _log_uniform(key, shape, minval, maxval):
+    assert 0 <= minval <= maxval
+    minval, maxval = map(jnp.log, (minval, maxval))
+    return jnp.exp(jax.random.uniform(key, shape, minval=minval, maxval=maxval))
+def setup_fn_randomize_damping_stiffness_factory(
+    prob_rigid: float,
+    all_imus_either_rigid_or_flex: bool,
+    imus_surely_rigid: list[str],
+):
+    assert 0 <= prob_rigid <= 1
+    assert prob_rigid != 1, "Use `imu_motion_artifacts`=False instead."
+    if prob_rigid == 0.0:
+        assert len(imus_surely_rigid) == 0
+    def stif_damp_rigid(key):
+        stif_sph = 200.0 * jnp.ones((3,))
+        stif_p3d = 2e4 * jnp.ones((3,))
+        stif = jnp.concatenate((stif_sph, stif_p3d))
+        return stif, stif * 0.2
+    def stif_damp_nonrigid(key):
+        keys = jax.random.split(key, 3)
+        stif_sph = _log_uniform(keys[0], (3,), _STIF_MIN_SPH, _STIF_MAX_SPH)
+        stif_p3d = _log_uniform(keys[1], (3,), _STIF_MIN_P3D, _STIF_MAX_P3D)
+        stif = jnp.concatenate((stif_sph, stif_p3d))
+        damp = _log_uniform(keys[2], (6,), _DAMP_MIN, _DAMP_MAX)
+        return stif, stif * damp
+    def setup_fn_randomize_damping_stiffness(key, sys: base.System) -> base.System:
+        link_damping = sys.link_damping
+        link_spring_stiffness = sys.link_spring_stiffness
+        idx_map = sys.idx_map("d")
+        imus = sys.findall_imus()
+        # initialize this RV because it might not get redrawn if
+        # `all_imus_either_rigid_or_flex` is set
+        key, consume = jax.random.split(key)
+        is_rigid = jax.random.bernoulli(consume, prob_rigid)
+        # this is only for the assertion used below
+        triggered_surely_rigid = []
+        for imu in imus:
+            # _imu has spherical joint and imu has p3d joint
+            slice = jnp.r_[idx_map[imu_reference_link_name(imu)], idx_map[imu]]
+            key, c1, c2 = jax.random.split(key, 3)
+            if prob_rigid > 0:
+                if imu in imus_surely_rigid:
+                    triggered_surely_rigid.append(imu)
+                    # logging.debug(f"IMU {imu} is surely rigid.")
+                    stif, damp = stif_damp_rigid(c2)
+                else:
+                    if not all_imus_either_rigid_or_flex:
+                        is_rigid = jax.random.bernoulli(c1, prob_rigid)
+                    stif, damp = jax.lax.cond(
+                        is_rigid, stif_damp_rigid, stif_damp_nonrigid, c2
+                    )
+            else:
+                stif, damp = stif_damp_nonrigid(c2)
+            link_spring_stiffness = link_spring_stiffness.at[slice].set(stif)
+            link_damping = link_damping.at[slice].set(damp)
+        assert len(imus_surely_rigid) == len(triggered_surely_rigid)
+        for imu_surely_rigid in imus_surely_rigid:
+            assert imu_surely_rigid in triggered_surely_rigid
+        return sys.replace(
+            link_damping=link_damping, link_spring_stiffness=link_spring_stiffness
+        )
+    return setup_fn_randomize_damping_stiffness
+def _match_q_x_between_sys(
+    sys_small: base.System,
+    q_large: jax.Array,
+    x_large: base.Transform,
+    sys_large: base.System,
+    q_large_skip: list[str],
+) -> tree_utils.PyTree:
+    assert q_large.ndim == 2
+    assert q_large.shape[1] == sys_large.q_size()
+    assert x_large.shape(1) == sys_large.num_links()
+    x_small_indices = []
+    q_small = []
+    q_idx_map = sys_large.idx_map("q")
+    def f(_, __, name: str):
+        x_small_indices.append(sys_large.name_to_idx(name))
+        # for the imu links the joint type was changed from spherical to frozen
+        # thus the q_idx_map has slices of length 4 but the `sys_small` has those
+        # imus but with frozen joint type and thus slices of length 0; so skip them
+        if name in q_large_skip:
+            return
+        q_small.append(q_large[:, q_idx_map[name]])
+    sys_small.scan(f, "l", sys_small.link_names)
+    x_small = tree_utils.tree_indices(x_large, jnp.array(x_small_indices), axis=1)
+    q_small = jnp.concatenate(q_small, axis=1)
+    return q_small, x_small
+class GeneratorTrafoHideInjectedBodies:
+    def __call__(self, gen):
+        def _gen(*args):
+            (X, y), (key, q, x, sys_x) = gen(*args)
+            # delete injected frames; then rename from `_imu` back to `imu`
+            imus = sys_x.findall_imus()
+            _imu2imu_map = {imu_reference_link_name(imu): imu for imu in imus}
+            sys = sys_x.delete_system(imus)
+            for _imu, imu in _imu2imu_map.items():
+                sys = sys.change_link_name(_imu, imu).change_joint_type(imu, "frozen")
+            # match q and x to `sys`; second axis is link axis
+            q, x = _match_q_x_between_sys(sys, q, x, sys_x, q_large_skip=imus)
+            return (X, y), (key, q, x, sys)
+        return _gen

ring/algorithms/generator/pd_control.py ADDED Viewed

@@ -0,0 +1,182 @@
+from types import SimpleNamespace
+from typing import Optional
+from flax import struct
+import jax
+import jax.numpy as jnp
+from ring import base
+from ring.algorithms import dynamics
+from ring.algorithms import jcalc
+@struct.dataclass
+class PDControllerState:
+    i: int
+    q_ref_as_dict: dict
+    qd_ref_as_dict: dict
+    P_gains: dict
+    D_gains: dict
+def _pd_control(P: jax.Array, D: Optional[jax.Array] = None):
+    """Computes tau using a PD controller. Returns a pair of (init, apply) functions.
+    NOTE: Gains around ~10_000 are good for spherical joints, everything else ~250-300
+    works just fine. Damping should be about 2500 for spherical joints, and
+    about 25 for everything else.
+    Args:
+        P: jax.Array of P gains. Shape: (sys_init.qd_size())
+        D: jax.Array of D gains. Shape: (sys_init.qd_size()) where `sys_init` is the
+            system that recorded the reference trajectory `q_ref`
+            If not given, then no D control is applied.
+    Returns: Pair of (init, apply) functions
+        init: (sys, q_ref) -> controller_state
+        apply: (controller_state, sys, state) -> controller_state, tau
+    Example:
+        >>> gains = jnp.array([250.0] * sys1.qd_size())
+        >>> controller = pd_control(gains, gains)
+        >>> q_ref = rcmg(sys1)
+        >>> cs = controller.init(sys1, q_ref)
+        >>> for t in range(1000):
+        >>>     cs, tau = controller.apply(cs, sys2, state)
+        >>>     state = dynamics.step(sys2, state, tau)
+    """
+    def init(sys: base.System, q_ref: jax.Array) -> dict:
+        assert sys.q_size() == q_ref.shape[1], f"q_ref.shape = {q_ref.shape}"
+        assert sys.qd_size() == P.size
+        if D is not None:
+            sys.qd_size() == D.size
+        q_ref_as_dict = {}
+        qd_ref_as_dict = {}
+        P_as_dict = {}
+        D_as_dict = {}
+        def f(_, __, q_ref_link, name, typ, P_link, D_link):
+            P_as_dict[name] = P_link
+            q_ref_link = q_ref_link.T
+            q_ref_as_dict[name] = q_ref_link
+            if D is not None:
+                qd_from_q = jcalc.get_joint_model(typ).qd_from_q
+                if qd_from_q is None:
+                    raise NotImplementedError(
+                        f"Please specify `JointModel.qd_from_q` for joint type `{typ}`"
+                    )
+                qd_ref_as_dict[name] = qd_from_q(q_ref_link, sys.dt)
+                D_as_dict[name] = D_link
+        sys.scan(
+            f,
+            "qlldd",
+            q_ref.T,
+            sys.link_names,
+            sys.link_types,
+            P,
+            D if D is not None else jnp.zeros((sys.qd_size(),)),
+        )
+        return PDControllerState(0, q_ref_as_dict, qd_ref_as_dict, P_as_dict, D_as_dict)
+    def apply(
+        controller_state: PDControllerState, sys: base.System, state: base.State
+    ) -> jax.Array:
+        taus = jnp.zeros((sys.qd_size()))
+        q_ref, qd_ref = jax.tree_map(
+            lambda arr: jax.lax.dynamic_index_in_dim(
+                arr, controller_state.i, keepdims=False
+            ),
+            (controller_state.q_ref_as_dict, controller_state.qd_ref_as_dict),
+        )
+        def f(_, idx_map, idx, name, typ, q_curr, qd_curr):
+            nonlocal taus
+            if name not in controller_state.q_ref_as_dict:
+                return
+            p_control_term = jcalc.get_joint_model(typ).p_control_term
+            if p_control_term is None:
+                raise NotImplementedError(
+                    f"Please specify `JointModel.p_control_term` for joint type `{typ}`"
+                )
+            P_term = p_control_term(q_curr, q_ref[name])
+            tau = P_term * controller_state.P_gains[name]
+            if name in controller_state.qd_ref_as_dict:
+                D_term = (qd_ref[name] - qd_curr) * controller_state.D_gains[name]
+                tau += D_term
+            taus = taus.at[idx_map["d"](idx)].set(tau)
+        sys.scan(
+            f,
+            "lllqd",
+            list(range(sys.num_links())),
+            sys.link_names,
+            sys.link_types,
+            state.q,
+            state.qd,
+        )
+        return controller_state.replace(i=controller_state.i + 1), taus
+    return SimpleNamespace(init=init, apply=apply)
+def _unroll_dynamics_pd_control(
+    sys: base.System,
+    q_ref: jax.Array,
+    P: jax.Array,
+    D: Optional[jax.Array] = None,
+    nograv: bool = False,
+    sys_q_ref: Optional[base.System] = None,
+    initial_sim_state_is_zeros: bool = False,
+    clip_taus: Optional[float] = None,
+):
+    assert q_ref.ndim == 2
+    if sys_q_ref is None:
+        sys_q_ref = sys
+    if nograv:
+        sys = sys.replace(gravity=sys.gravity * 0.0)
+    if initial_sim_state_is_zeros:
+        state = base.State.create(sys)
+    else:
+        state = _initial_q_is_q_ref(sys, sys_q_ref, q_ref[0])
+    controller = _pd_control(P, D)
+    cs = controller.init(sys_q_ref, q_ref)
+    def step(carry, _):
+        state, cs = carry
+        cs, taus = controller.apply(cs, sys, state)
+        if clip_taus is not None:
+            assert clip_taus > 0.0
+            taus = jnp.clip(taus, -clip_taus, clip_taus)
+        state = dynamics.step(sys, state, taus)
+        carry = (state, cs)
+        return carry, state
+    states = jax.lax.scan(step, (state, cs), None, length=q_ref.shape[0])[1]
+    return states
+def _initial_q_is_q_ref(sys: base.System, sys_q_ref: base.System, q_ref):
+    # you can not preallocate q using zeros because of quaternions..
+    q = base.State.create(sys).q
+    sys_q_map = sys.idx_map("q")
+    def f(_, __, name, q_ref_link):
+        nonlocal q
+        q = q.at[sys_q_map[name]].set(q_ref_link)
+    sys_q_ref.scan(f, "lq", sys_q_ref.link_names, q_ref)
+    return base.State.create(sys, q=q)

ring/algorithms/generator/randomize.py ADDED Viewed

@@ -0,0 +1,119 @@
+"""Randomization by modifying System and MotionConfig objects before building
+generator."""
+from dataclasses import replace
+import itertools
+from typing import Optional
+import warnings
+import jax.numpy as jnp
+from ring import base
+from ring.algorithms import jcalc
+from ring.algorithms.generator import types
+def _find_children(lam: list[int], body: int) -> list[int]:
+    children = []
+    def _children(body: int) -> None:
+        for i in range(len(lam)):
+            if lam[i] == body:
+                children.append(i)
+                _children(i)
+    _children(body)
+    return children
+def _find_root_of_subsys_that_contains_body(sys: base.System, body: str) -> str:
+    body_i = sys.name_to_idx(body)
+    for i, p in enumerate(sys.link_parents):
+        if p == -1:
+            if body_i == i or body_i in _find_children(sys.link_parents, i):
+                return sys.idx_to_name(i)
+def _assign_anchors_to_subsys(sys: base.System, anchors: list[str]) -> list[list[str]]:
+    anchors_per_subsys = []
+    for i, p in enumerate(sys.link_parents):
+        if p == -1:
+            link_idxs_subsys = [i] + _find_children(sys.link_parents, i)
+            link_names_subsys = [sys.idx_to_name(i) for i in link_idxs_subsys]
+            anchors_this_subsys = [
+                name for name in anchors if name in link_names_subsys
+            ]
+            if len(anchors_this_subsys) == 0:
+                anchors_this_subsys = [sys.idx_to_name(i)]
+            anchors_per_subsys.append(anchors_this_subsys)
+    return anchors_per_subsys
+def _morph_extract_subsys(sys: base.System, anchor: str):
+    root = _find_root_of_subsys_that_contains_body(sys, anchor)
+    roots = sys.findall_bodies_to_world(names=True)
+    subsys = sys.delete_system(list(set(roots) - set([root])))
+    return subsys.morph_system(new_anchor=anchor)
+def randomize_anchors(
+    sys: base.System, anchors: Optional[list[str]] = None
+) -> list[base.System]:
+    if anchors is None:
+        anchors = sys.findall_segments()
+    anchors = _assign_anchors_to_subsys(sys, anchors)
+    syss = []
+    for anchors_subsys in itertools.product(*anchors):
+        sys_mod = _morph_extract_subsys(sys, anchors_subsys[0])
+        for anchor_subsys in anchors_subsys[1:]:
+            sys_mod = sys_mod.inject_system(_morph_extract_subsys(sys, anchor_subsys))
+        syss.append(sys_mod)
+    return syss
+_WARN_HZ_Threshold: float = 40.0
+def randomize_hz(
+    sys: list[base.System],
+    configs: list[jcalc.MotionConfig],
+    sampling_rates: list[float],
+) -> tuple[list[base.System], list[jcalc.MotionConfig]]:
+    Ts = [c.T for c in configs]
+    assert len(set(Ts)), f"Time length between configs does not agree {Ts}"
+    T_global = Ts[0]
+    for hz in sampling_rates:
+        if hz < _WARN_HZ_Threshold:
+            warnings.warn(
+                "The sampling rate {hz} is below the warning threshold of "
+                f"{_WARN_HZ_Threshold}. This might lead to NaNs."
+            )
+    sys_out, configs_out = [], []
+    for _sys in sys:
+        for _config in configs:
+            for hz in sampling_rates:
+                dt = 1 / hz
+                T = (T_global / _sys.dt) * dt
+                sys_out.append(_sys.replace(dt=dt))
+                configs_out.append(replace(_config, T=T))
+    return sys_out, configs_out
+def randomize_hz_finalize_fn_factory(finalize_fn_user: types.FINALIZE_FN | None):
+    def finalize_fn(key, q, x, sys: base.System):
+        X, y = {}, {}
+        if finalize_fn_user is not None:
+            X, y = finalize_fn_user(key, q, x, sys)
+        assert "dt" not in X
+        X["dt"] = jnp.array([sys.dt], dtype=jnp.float32)
+        return X, y
+    return finalize_fn