imt-ring 1.6.20__tar.gz → 1.6.22__tar.gz

{imt_ring-1.6.20 → imt_ring-1.6.22}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: imt-ring
-Version: 1.6.20
+Version: 1.6.22
 Summary: RING: Recurrent Inertial Graph-based Estimator
 Author-email: Simon Bachhuber <simon.bachhuber@fau.de>
 Project-URL: Homepage, https://github.com/SimiPixel/ring

{imt_ring-1.6.20 → imt_ring-1.6.22}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "imt-ring"
-version = "1.6.20"
+version = "1.6.22"
 authors = [
   { name="Simon Bachhuber", email="simon.bachhuber@fau.de" },
 ]

{imt_ring-1.6.20 → imt_ring-1.6.22}/src/imt_ring.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: imt-ring
-Version: 1.6.20
+Version: 1.6.22
 Summary: RING: Recurrent Inertial Graph-based Estimator
 Author-email: Simon Bachhuber <simon.bachhuber@fau.de>
 Project-URL: Homepage, https://github.com/SimiPixel/ring

{imt_ring-1.6.20 → imt_ring-1.6.22}/src/ring/algorithms/_random.py

@@ -199,7 +199,7 @@ def random_position_over_time(
     POS = jnp.zeros((int(T // t_min) + 1, 2))
     POS = POS.at[0, 1].set(POS_0)
 
-    val_outer = (1, 0.0, 0.0, 0.0, 0.0, key, POS)
+    val_outer = (1, 0.0, 0.0, POS_0, POS_0, key, POS)
     end, *_, consume, POS = jax.lax.while_loop(cond_fn_outer, body_fn_outer, val_outer)
     POS = jnp.where(
         (jnp.arange(len(POS)) < end)[:, None],

{imt_ring-1.6.20 → imt_ring-1.6.22}/src/ring/algorithms/dynamics.py

@@ -190,11 +190,11 @@ def _spring_force(sys: base.System, q: jax.Array):
 
     def _calc_spring_force_per_link(_, __, q, zeropoint, typ):
         # cor is (free, p3d) stacked; free is (spherical, p3d) stacked
-        if typ in ["free", "cor"]:
+        if base.System.joint_type_is_free_or_cor(typ):
             quat_force = _quaternion_spring_force(zeropoint[:4], q[:4])
             pos_force = zeropoint[4:] - q[4:]
             q_spring_force_link = jnp.concatenate((quat_force, pos_force))
-        elif typ == "spherical":
+        elif base.System.joint_type_is_spherical(typ):
             q_spring_force_link = _quaternion_spring_force(zeropoint, q)
         else:
             q_spring_force_link = zeropoint - q
@@ -268,11 +268,11 @@ def _semi_implicit_euler_integration(
     q_next = []
 
     def q_integrate(_, __, q, qd, typ):
-        if typ in ["free", "cor"]:
+        if sys.joint_type_is_free_or_cor(typ):
             quat_next = _strapdown_integration(q[:4], qd[:3], sys.dt)
             pos_next = q[4:] + qd[3:] * sys.dt
             q_next_i = jnp.concatenate((quat_next, pos_next))
-        elif typ == "spherical":
+        elif sys.joint_type_is_spherical(typ):
             quat_next = _strapdown_integration(q, qd, sys.dt)
             q_next_i = quat_next
         else:

{imt_ring-1.6.20 → imt_ring-1.6.22}/src/ring/algorithms/generator/base.py

@@ -1,3 +1,4 @@
+from dataclasses import replace
 from functools import partial
 import random
 from typing import Callable, Optional
@@ -34,6 +35,8 @@ class RCMG:
         add_y_relpose: bool = False,
         add_y_rootincl: bool = False,
         add_y_rootincl_kwargs: dict = dict(),
+        add_y_rootfull: bool = False,
+        add_y_rootfull_kwargs: dict = dict(),
         sys_ml: Optional[base.System] = None,
         randomize_positions: bool = False,
         randomize_motion_artifacts: bool = False,
@@ -73,6 +76,8 @@ class RCMG:
                     add_y_relpose=add_y_relpose,
                     add_y_rootincl=add_y_rootincl,
                     add_y_rootincl_kwargs=add_y_rootincl_kwargs,
+                    add_y_rootfull=add_y_rootfull,
+                    add_y_rootfull_kwargs=add_y_rootfull_kwargs,
                     sys_ml=sys_ml,
                     randomize_positions=randomize_positions,
                     randomize_motion_artifacts=randomize_motion_artifacts,
@@ -279,6 +284,8 @@ def _build_mconfig_batched_generator(
     add_y_relpose: bool,
     add_y_rootincl: bool,
     add_y_rootincl_kwargs: dict,
+    add_y_rootfull: bool,
+    add_y_rootfull_kwargs: dict,
     sys_ml: base.System,
     randomize_positions: bool,
     randomize_motion_artifacts: bool,
@@ -365,7 +372,11 @@ def _build_mconfig_batched_generator(
         if add_y_relpose:
             pipe.append(finalize_fns.RelPose(sys_noimu))
         if add_y_rootincl:
+            assert not add_y_rootfull
             pipe.append(finalize_fns.RootIncl(sys_noimu, **add_y_rootincl_kwargs))
+        if add_y_rootfull:
+            assert not add_y_rootincl
+            pipe.append(finalize_fns.RootFull(sys_noimu, **add_y_rootfull_kwargs))
         if use_link_number_in_Xy:
             pipe.append(finalize_fns.Names2Indices(sys_noimu))
 
@@ -436,7 +447,15 @@ def draw_random_q(
         draw_fn = jcalc.get_joint_model(link_type).rcmg_draw_fn
         if draw_fn is None:
             raise Exception(f"The joint type {link_type} has no draw fn specified.")
-        q_link = draw_fn(config, key_t, key_value, sys.dt, N, joint_params)
+
+        if link_type in config.joint_type_specific_overwrites:
+            _config = replace(
+                config, **config.joint_type_specific_overwrites[link_type]
+            )
+        else:
+            _config = config
+
+        q_link = draw_fn(_config, key_t, key_value, sys.dt, N, joint_params)
         # even revolute and prismatic joints must be 2d arrays
         q_link = q_link if q_link.ndim == 2 else q_link[:, None]
         q_list.append(q_link)

{imt_ring-1.6.20 → imt_ring-1.6.22}/src/ring/algorithms/generator/finalize_fns.py

@@ -88,6 +88,18 @@ class RootIncl:
         return (X, y), (key, q, x, sys_x)
 
 
+class RootFull:
+    def __init__(self, sys: base.System, **kwargs):
+        self.sys = sys
+        self.kwargs = kwargs
+
+    def __call__(self, Xy, extras):
+        (X, y), (key, q, x, sys_x) = Xy, extras
+        y_root_incl = sensors.root_full(self.sys, x, sys_x, **self.kwargs)
+        y = utils.dict_union(y, y_root_incl)
+        return (X, y), (key, q, x, sys_x)
+
+
 _default_imu_kwargs = dict(
     noisy=True,
     low_pass_filter_pos_f_cutoff=13.5,

{imt_ring-1.6.20 → imt_ring-1.6.22}/src/ring/algorithms/jcalc.py

@@ -40,6 +40,12 @@ class MotionConfig:
     dpos_max: float | TimeDependentFloat = 0.7
     pos_min: float | TimeDependentFloat = -2.5
     pos_max: float | TimeDependentFloat = +2.5
+    pos_min_p3d_x: float | TimeDependentFloat = -2.5
+    pos_max_p3d_x: float | TimeDependentFloat = +2.5
+    pos_min_p3d_y: float | TimeDependentFloat = -2.5
+    pos_max_p3d_y: float | TimeDependentFloat = +2.5
+    pos_min_p3d_z: float | TimeDependentFloat = -2.5
+    pos_max_p3d_z: float | TimeDependentFloat = +2.5
 
     # used by both `random_angle_*` and `random_pos_*`
     # only used if `randomized_interpolation` is set
@@ -59,6 +65,12 @@ class MotionConfig:
     ang0_max: float = jnp.pi
     pos0_min: float = 0.0
     pos0_max: float = 0.0
+    pos0_min_p3d_x: float = 0.0
+    pos0_max_p3d_x: float = 0.0
+    pos0_min_p3d_y: float = 0.0
+    pos0_max_p3d_y: float = 0.0
+    pos0_min_p3d_z: float = 0.0
+    pos0_max_p3d_z: float = 0.0
 
     # cor (center of rotation) custom fields
     cor_t_min: float = 0.2
@@ -67,6 +79,14 @@ class MotionConfig:
     cor_dpos_max: float | TimeDependentFloat = 0.5
     cor_pos_min: float | TimeDependentFloat = -0.4
     cor_pos_max: float | TimeDependentFloat = 0.4
+    cor_pos0_min: float = 0.0
+    cor_pos0_max: float = 0.0
+
+    # specify changes for this motionconfig and for specific joint types
+    # map of `link_types` -> dictionary of changes
+    joint_type_specific_overwrites: dict[str, dict[str, Any]] = field(
+        default_factory=lambda: dict()
+    )
 
     def is_feasible(self) -> bool:
         return _is_feasible_config1(self)
@@ -92,6 +112,9 @@ class MotionConfig:
     def overwrite_for_joint_type(joint_type: str, **changes) -> None:
         """Changes values of the `MotionConfig` used by the draw_fn for only a specific
         joint.
+        !!! Note
+            This applies these changes to *all* MotionConfigs for this joint type!
+            This takes precedence *over* `Motionconfig.joint_type_specific_overwrites`!
         """
         previous_changes = _overwrite_for_joint_type_changes[joint_type]
         for change in changes:
@@ -113,6 +136,56 @@ class MotionConfig:
             overwrite=True,
         )
 
+    @staticmethod
+    def overwrite_for_subsystem(
+        sys: base.System, link_name: str, **changes
+    ) -> base.System:
+        """Modifies motionconfig of all joints in subsystem with root `link_name`.
+        Note that if the subsystem contains a free joint then the jointtype will
+        will be re-named to `free_<link_name>`, then the RCMG flag `cor` will
+        potentially not work as expected because it searches for all joints of
+        type `free` to replace with `cor`. The workaround here is to change the
+        type already from `free` to `cor in the xml file.
+        This takes precedence *over* `Motionconfig.joint_type_specific_overwrites`!
+
+        Args:
+            sys (base.System): System object that gets updated
+            link_name (str): Root node of subsystem
+            changes: Changes to apply to the motionconfig
+
+        Return:
+            base.System: Updated system with new jointtypes
+        """
+        from ring.algorithms.generator.finalize_fns import _P_gains
+
+        # all bodies in the subsystem
+        bodies = sys.findall_bodies_subsystem(link_name) + [sys.name_to_idx(link_name)]
+
+        jts_subsys = set([sys.link_types[i] for i in bodies]) - set(["frozen"])
+        postfix = "_" + link_name
+        # create new joint types with updated motionconfig
+        for typ in jts_subsys:
+            register_new_joint_type(
+                typ + postfix,
+                get_joint_model(typ),
+                base.Q_WIDTHS[typ],
+                base.QD_WIDTHS[typ],
+            )
+            MotionConfig.overwrite_for_joint_type(typ + postfix, **changes)
+            _P_gains[typ + postfix] = _P_gains[typ]
+
+        # rename all jointtypes
+        new_link_types = [
+            (
+                sys.link_types[i] + postfix
+                if (i in bodies and sys.link_types[i] != "frozen")
+                else sys.link_types[i]
+            )
+            for i in range(sys.num_links())
+        ]
+        sys = sys.replace(link_types=new_link_types)
+        return sys
+
     @staticmethod
     def from_register(name: str) -> "MotionConfig":
         return _registered_motion_configs[name]
@@ -221,6 +294,37 @@ _registered_motion_configs = {
 }
 
 
+def _joint_specific_overwrites_free_cor(
+    id: str, dang: float, dpos: float
+) -> MotionConfig:
+    changes = dict(
+        dang_max_free_spherical=dang,
+        dpos_max=dpos,
+        cor_dpos_max=dpos,
+        t_min=1.5,
+        t_max=15.0,
+    )
+    return replace(
+        _registered_motion_configs[id],
+        joint_type_specific_overwrites=dict(free=changes, cor=changes),
+    )
+
+
+_registered_motion_configs.update(
+    {
+        f"{id}-S": _joint_specific_overwrites_free_cor(id, 0.2, 0.1)
+        for id in ["expSlow", "expFast", "hinUndHer", "standard"]
+    }
+)
+_registered_motion_configs.update(
+    {
+        f"{id}-S+": _joint_specific_overwrites_free_cor(id, 0.1, 0.05)
+        for id in ["expSlow", "expFast", "hinUndHer", "standard"]
+    }
+)
+del _joint_specific_overwrites_free_cor
+
+
 def _is_feasible_config1(c: MotionConfig) -> bool:
     t_min, t_max = c.t_min, _to_float(c.t_max, 0.0)
 
@@ -254,8 +358,26 @@ def _is_feasible_config1(c: MotionConfig) -> bool:
     cond2 = inside_box_checks(
         _to_float(c.pos_min, 0.0), _to_float(c.pos_max, 0.0), c.pos0_min, c.pos0_max
     )
+    cond3 = inside_box_checks(
+        _to_float(c.pos_min_p3d_x, 0.0),
+        _to_float(c.pos_max_p3d_x, 0.0),
+        c.pos0_min_p3d_x,
+        c.pos0_max_p3d_x,
+    )
+    cond4 = inside_box_checks(
+        _to_float(c.pos_min_p3d_y, 0.0),
+        _to_float(c.pos_max_p3d_y, 0.0),
+        c.pos0_min_p3d_y,
+        c.pos0_max_p3d_y,
+    )
+    cond5 = inside_box_checks(
+        _to_float(c.pos_min_p3d_z, 0.0),
+        _to_float(c.pos_max_p3d_z, 0.0),
+        c.pos0_min_p3d_z,
+        c.pos0_max_p3d_z,
+    )
 
-    return cond1 and cond2
+    return cond1 and cond2 and cond3 and cond4 and cond5
 
 
 def _find_interval(t: jax.Array, boundaries: jax.Array):
@@ -504,7 +626,11 @@ def _draw_pxyz(
     cor: bool = False,
 ) -> jax.Array:
     key_value, consume = jax.random.split(key_value)
-    POS_0 = jax.random.uniform(consume, minval=config.pos0_min, maxval=config.pos0_max)
+    POS_0 = jax.random.uniform(
+        consume,
+        minval=config.cor_pos0_min if cor else config.pos0_min,
+        maxval=config.cor_pos0_max if cor else config.pos0_max,
+    )
     max_iter = 100
     return _random.random_position_over_time(
         key_value,
@@ -590,10 +716,27 @@ def _draw_p3d_and_cor(
     __: jax.Array,
     cor: bool,
 ) -> jax.Array:
-    pos = jax.vmap(lambda key: _draw_pxyz(config, None, key, dt, N, None, cor))(
-        jax.random.split(key_value, 3)
-    )
-    return pos.T
+    keys = jax.random.split(key_value, 3)
+
+    def draw(key, xyz: str):
+        return _draw_pxyz(
+            replace(
+                config,
+                pos_min=getattr(config, f"pos_min_p3d_{xyz}"),
+                pos_max=getattr(config, f"pos_max_p3d_{xyz}"),
+                pos0_min=getattr(config, f"pos0_min_p3d_{xyz}"),
+                pos0_max=getattr(config, f"pos0_max_p3d_{xyz}"),
+            ),
+            None,
+            key,
+            dt,
+            N,
+            None,
+            cor,
+        )[:, None]
+
+    px, py, pz = draw(keys[0], "x"), draw(keys[1], "y"), draw(keys[2], "z")
+    return jnp.concat((px, py, pz), axis=-1)
 
 
 def _draw_p3d(

{imt_ring-1.6.20 → imt_ring-1.6.22}/src/ring/base.py

@@ -7,6 +7,7 @@ from jax.core import Tracer
 import jax.numpy as jnp
 from jax.tree_util import tree_map
 import numpy as np
+import tree
 import tree_utils as tu
 
 import ring
@@ -590,6 +591,34 @@ class System(_Base):
 
         return sys
 
+    @staticmethod
+    def joint_type_simplification(typ: str) -> str:
+        if typ[:4] == "free":
+            if typ == "free_2d":
+                return "free_2d"
+            else:
+                return "free"
+        elif typ[:3] == "cor":
+            return "cor"
+        elif typ[:9] == "spherical":
+            return "spherical"
+        else:
+            return typ
+
+    @staticmethod
+    def joint_type_is_free_or_cor(typ: str) -> bool:
+        return System.joint_type_simplification(typ) in ["free", "cor"]
+
+    @staticmethod
+    def joint_type_is_spherical(typ: str) -> bool:
+        return System.joint_type_simplification(typ) == "spherical"
+
+    @staticmethod
+    def joint_type_is_free_or_cor_or_spherical(typ: str) -> bool:
+        return System.joint_type_is_free_or_cor(typ) or System.joint_type_is_spherical(
+            typ
+        )
+
     def findall_imus(self, names: bool = True) -> list[str] | list[int]:
         bodies = [name for name in self.link_names if name[:3] == "imu"]
         return bodies if names else [self.name_to_idx(n) for n in bodies]
@@ -618,10 +647,20 @@ class System(_Base):
         return self._bodies_indices_to_bodies_name(bodies) if names else bodies
 
     def children(self, name: str, names: bool = False) -> list[int] | list[str]:
+        "List all direct children of body, does not include body itself"
         p = self.name_to_idx(name)
         bodies = [i for i in range(self.num_links()) if self.link_parents[i] == p]
         return bodies if (not names) else [self.idx_to_name(i) for i in bodies]
 
+    def findall_bodies_subsystem(
+        self, name: str, names: bool = False
+    ) -> list[int] | list[str]:
+        "List all children and children's children; does not include body itself"
+        children = self.children(name, names=True)
+        grandchildren = [self.findall_bodies_subsystem(n, names=True) for n in children]
+        bodies = tree.flatten([children, grandchildren])
+        return bodies if names else [self.name_to_idx(n) for n in bodies]
+
     def scan(self, f: Callable, in_types: str, *args, reverse: bool = False):
         """Scan `f` along each link in system whilst carrying along state.
 
@@ -889,7 +928,9 @@ def _parse_system(sys: System) -> System:
         assert d.size == a.size == s.size == qd_size, error_msg
         assert z.size == q_size, error_msg
 
-        if typ in ["spherical", "free", "cor"] and not isinstance(z, Tracer):
+        if System.joint_type_is_free_or_cor_or_spherical(typ) and not isinstance(
+            z, Tracer
+        ):
             assert jnp.allclose(
                 jnp.linalg.norm(z[:4]), 1.0
             ), f"not unit quat for link `{name}` of typ `{typ}` in model"
@@ -1030,7 +1071,7 @@ class State(_Base):
             def replace_by_unit_quat(_, idx_map, link_typ, link_idx):
                 nonlocal q
 
-                if link_typ in ["free", "cor", "spherical"]:
+                if sys.joint_type_is_free_or_cor_or_spherical(link_typ):
                     q_idxs_link = idx_map["q"](link_idx)
                     q = q.at[q_idxs_link.start].set(1.0)
 

{imt_ring-1.6.20 → imt_ring-1.6.22}/src/ring/io/xml/abstract.py

@@ -3,6 +3,7 @@ from typing import Tuple, TypeVar
 import jax
 import jax.numpy as jnp
 import numpy as np
+
 from ring import base
 
 T = TypeVar("T")
@@ -17,7 +18,7 @@ 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"]:
+    if base.System.joint_type_is_free_or_cor_or_spherical(link_typ):
         # zeropoint then is unit quaternion and not zeros
         zeropoint = zeropoint.at[0].set(1.0)
     return zeropoint

{imt_ring-1.6.20 → imt_ring-1.6.22}/src/ring/ml/train.py

@@ -167,7 +167,10 @@ def train_fn(
         tbp=tbp,
     )
 
-    default_callbacks = []
+    # always log, because we also want `i_epsiode` to be logged in wandb
+    default_callbacks = [
+        ml_callbacks.LogEpisodeTrainingLoopCallback(callback_kill_after_episode)
+    ]
     if metrices is not None:
         eval_fn = _build_eval_fn(metrices, filter, link_names)
         default_callbacks.append(_DefaultEvalFnCallback(eval_fn))
@@ -192,11 +195,6 @@ def train_fn(
     if callback_kill_if_nan:
         default_callbacks.append(ml_callbacks.NanKillRunCallback())
 
-    # always log, because we also want `i_epsiode` to be logged in wandb
-    default_callbacks.append(
-        ml_callbacks.LogEpisodeTrainingLoopCallback(callback_kill_after_episode)
-    )
-
     if callback_kill_after_seconds is not None:
         default_callbacks.append(
             ml_callbacks.TimingKillRunCallback(callback_kill_after_seconds)