imt-ring 1.6.20__py3-none-any.whl → 1.6.22__py3-none-any.whl

{imt_ring-1.6.20.dist-info → imt_ring-1.6.22.dist-info}/METADATA

@@ -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.dist-info → imt_ring-1.6.22.dist-info}/RECORD

@@ -1,12 +1,12 @@
 ring/__init__.py,sha256=ncBRdHvge6uqpzFyk8_HUQNx4kZxENpVXTJTEK_SjCg,5216
 ring/algebra.py,sha256=F0GwbP8LQP5qGVkoMUYJmkp9Hn2nKAVIkCVYDEjNjGU,3128
-ring/base.py,sha256=Ystn1EjTyOXBhVm5koroV_YPUYtFxrteJLd-XR3kEL8,33840
+ring/base.py,sha256=yPdbPywwDllCRsJEbnLW4s9Z-bBD8qdxpEDYV3pCLP8,35296
 ring/maths.py,sha256=qPHH6TpHCK3TgExI98gNEySoSRKOwteN9McUlyUFipI,12207
 ring/spatial.py,sha256=nmZ-UhRanhyM34bez8uCS4wMwaKqLkuEbgKGP5XNH60,2351
 ring/algorithms/__init__.py,sha256=IiK9EN5Xgs3dB075-A-H-Yad0Z7vzvKIJF2g6X_-C_8,1224
-ring/algorithms/_random.py,sha256=fc26yEQjSjtf0NluZ41CyeGIRci0ldrRlThueHR9H7U,14007
-ring/algorithms/dynamics.py,sha256=GOedL1STj6oXcXgMA7dB4PabvCQxPBbirJQhXBRuKqE,10929
-ring/algorithms/jcalc.py,sha256=kQgrRE0XoUBrcdeFbUw_xKnL4m1P2G1Q1m4n7BX2yDk,30064
+ring/algorithms/_random.py,sha256=6QwV7EAE2ckDi7rSGXT1XcPhR5R-NU-J8HVa5DoPBSQ,14011
+ring/algorithms/dynamics.py,sha256=dpe-F3Yq4sY2dY6DQW3v7TnPLRdxdkePtdbGPQIrijg,10997
+ring/algorithms/jcalc.py,sha256=pniA8uZutN6H_4sP_YxHbGfW-mWEEUimHGkwLC0fgms,35049
 ring/algorithms/kinematics.py,sha256=DOboHI517Vx0pRJUFZtZPmK_qFaiKiQe-37B-M0aC-c,7422
 ring/algorithms/sensors.py,sha256=0xOzdQIc1kBF0CkoPXWWCx3MmV4SG3wj7knVnnMWq9M,18124
 ring/algorithms/custom_joints/__init__.py,sha256=3pQ-Is_HBTQDkzESCNg9VfoP8wvseWmooryG8ERnu_A,366
@@ -15,9 +15,9 @@ ring/algorithms/custom_joints/rr_joint.py,sha256=jnRtjtOCALMaq2_0bcu2d7qgfQ6etXp
 ring/algorithms/custom_joints/rsaddle_joint.py,sha256=QoMo6NXdYgA9JygSzBvr0eCdd3qKhUgCrGPNO2Qdxko,1200
 ring/algorithms/custom_joints/suntay.py,sha256=tOEGM304XciHO4pmvxr4faA4xXVO4N2HlPdFmXKbcrw,16726
 ring/algorithms/generator/__init__.py,sha256=bF-CW3x2x-o6KWESKy-DuxzZPh3UNSjJb_MaAcSHGsQ,277
-ring/algorithms/generator/base.py,sha256=vxUdA0ZeSNH3SOanL51qVRvCiJrmsWQyQX0g2fdm3Rg,15825
+ring/algorithms/generator/base.py,sha256=nFVUcFl7AILnyjuXf_YQJxI12MB4TzJ1SRFREKzqv8Q,16531
 ring/algorithms/generator/batch.py,sha256=9yFxVv11hij-fJXGPxA3zEh1bE2_jrZk0R7kyGaiM5c,2551
-ring/algorithms/generator/finalize_fns.py,sha256=EZ5p7fuZu0Zd0rHJzCVg3vy3U9ysny6TfQfolIGPERc,10029
+ring/algorithms/generator/finalize_fns.py,sha256=nY2RKiLbHriTkdec94lc4UGSZKd0v547MDNn4dr8I3E,10398
 ring/algorithms/generator/motion_artifacts.py,sha256=2VJbldVDbI3PSyboshIbtYvSAKzBBwGV7cQfYjqvluM,9167
 ring/algorithms/generator/pd_control.py,sha256=XJ_Gd5AkIRh-jBrMfQyMXjVwhx2gCNHznjzFbmAwhZs,5767
 ring/algorithms/generator/setup_fns.py,sha256=MFz3czHBeWs1Zk1A8O02CyQpQ-NCyW9PMpbqmKit6es,1455
@@ -46,7 +46,7 @@ ring/io/examples/exclude/standard_sys_rr_imp.xml,sha256=1K8aLe4n97qFMQaytXdHVHDf
 ring/io/examples/test_morph_system/four_seg_seg1.xml,sha256=XJvGtEnvedejs_OmCVfQULWJNK8MLDQQ3raqPNRCJbA,1283
 ring/io/examples/test_morph_system/four_seg_seg3.xml,sha256=HktN7_a_Ly3YflWit5W-WncxApWGMORAGnRXyMEqnoA,1265
 ring/io/xml/__init__.py,sha256=-3k6ffvFyc4zm0oTyVz3ez-o3Lb9bPp2sjwSub_K1AA,242
-ring/io/xml/abstract.py,sha256=ojsXgz15J4pI1FThUIZbB_Iw1wmR9cOHt5Thmcbih4I,9721
+ring/io/xml/abstract.py,sha256=8Q2ebnUYLmuS9HJAQwDVrDTrRfD5z4G5RAB7MW8Oa60,9742
 ring/io/xml/from_xml.py,sha256=8b44sPVWgoY8JGJZLpJ8M_eLfcfu3IsMtBzSytPTPmw,9234
 ring/io/xml/test_from_xml.py,sha256=bckVrVVmEhCwujd_OF9FGYnX3zU3BgztpqGxxmd0htM,1562
 ring/io/xml/test_to_xml.py,sha256=NGn4VSiFdwhYN5YTBduWMiY9B5dwtxZhCQAR_PXeqKU,946
@@ -58,7 +58,7 @@ ring/ml/ml_utils.py,sha256=xqy9BnLy8IKVqkFS9mlZsGJXSbThI9zZxZ5rhl8LSI8,7144
 ring/ml/optimizer.py,sha256=TZF0_LmnewzmGVso-zIQJtpWguUW0fW3HeRpIdG_qoI,4763
 ring/ml/ringnet.py,sha256=mef7jyN2QcApJmQGH3HYZyTV-00q8YpsYOKhW0-ku1k,8973
 ring/ml/rnno_v1.py,sha256=2qE08OIvTJ5PvSxKpYGzGSrvEImWrdAT_qslZ7jP5tA,1372
-ring/ml/train.py,sha256=huUfMK6eotS6BRrQKoZ-AUG0um3jlqpfQFZNJT8LKiE,10854
+ring/ml/train.py,sha256=XuUUB0NhvByGtZDtS_weyp-TKPG9ErnKixS4NqB8q6M,10822
 ring/ml/training_loop.py,sha256=CEokvPQuuk_WCd-J60ZDodJYcPVvyxLfgXDr_DnbzRI,3359
 ring/ml/params/0x13e3518065c21cd8.pickle,sha256=Zh2k1zK-TNxJl5F7nyTeQ9001qqRE_dfvaq1HWV287A,9355838
 ring/ml/params/0x1d76628065a71e0f.pickle,sha256=YTNVuvfw-nCRD9BH1PZYcR9uCFpNWDhw8Lc50eDn_EE,9351038
@@ -86,7 +86,7 @@ ring/utils/randomize_sys.py,sha256=G_vBIo0OwQkXL2u0djwbaoaeb02C4LQCTNNloOYIU2M,3
 ring/utils/utils.py,sha256=tJaWXLGOTwkxJQj2l23dX97wO3aZYhM2qd7eNuMRs84,6907
 ring/utils/register_gym_envs/__init__.py,sha256=PtPIRBQJ16339xZ9G9VpvqrvcGbQ_Pk_SUz4tQPa9nQ,94
 ring/utils/register_gym_envs/saddle.py,sha256=tA5CyW_akSXyDm0xJ83CtOrUMVElH0f9vZtEDDJQalI,4422
-imt_ring-1.6.20.dist-info/METADATA,sha256=uOgzTpTz7MgDEIGhL7onaD_r64DlL6ALSnaT5t8u3HY,4097
-imt_ring-1.6.20.dist-info/WHEEL,sha256=a7TGlA-5DaHMRrarXjVbQagU3Man_dCnGIWMJr5kRWo,91
-imt_ring-1.6.20.dist-info/top_level.txt,sha256=EiT790-lAyi8iwTzJArH3f2k77rwhDn00q-4PlmvDQo,5
-imt_ring-1.6.20.dist-info/RECORD,,
+imt_ring-1.6.22.dist-info/METADATA,sha256=EHw8XK0JO8d6lQ-PiknLrM7V5ATyLkniMzBFydE6NA4,4097
+imt_ring-1.6.22.dist-info/WHEEL,sha256=R06PA3UVYHThwHvxuRWMqaGcr-PuniXahwjmQRFMEkY,91
+imt_ring-1.6.22.dist-info/top_level.txt,sha256=EiT790-lAyi8iwTzJArH3f2k77rwhDn00q-4PlmvDQo,5
+imt_ring-1.6.22.dist-info/RECORD,,

{imt_ring-1.6.20.dist-info → imt_ring-1.6.22.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (75.4.0)
+Generator: setuptools (75.5.0)
 Root-Is-Purelib: true
 Tag: py3-none-any
 

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],

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:

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)

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,

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(

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)
 

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

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)