imt-ring 1.6.15__py3-none-any.whl → 1.6.17__py3-none-any.whl

{imt_ring-1.6.15.dist-info → imt_ring-1.6.17.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: imt-ring
-Version: 1.6.15
+Version: 1.6.17
 Summary: RING: Recurrent Inertial Graph-based Estimator
 Author-email: Simon Bachhuber <simon.bachhuber@fau.de>
 Project-URL: Homepage, https://github.com/SimiPixel/ring
@@ -32,11 +32,11 @@ Requires-Dist: pytest-xdist ; extra == 'dev'
 Requires-Dist: nbmake ; extra == 'dev'
 
 <p align="center">
-<img src="https://raw.githubusercontent.com/SimiPixel/ring/main/docs/img/icon.svg" height="200" />
+<img src="https://raw.githubusercontent.com/simon-bachhuber/ring/main/docs/img/icon.svg" height="200" />
 </p>
 
 # Recurrent Inertial Graph-based Estimator (RING)
-<img src="https://raw.githubusercontent.com/SimiPixel/ring/main/docs/img/coverage_badge.svg" height="20" />
+<img src="https://raw.githubusercontent.com/simon-bachhuber/ring/main/docs/img/coverage_badge.svg" height="20" />
 
 ## Installation
 

{imt_ring-1.6.15.dist-info → imt_ring-1.6.17.dist-info}/RECORD

@@ -1,4 +1,4 @@
-ring/__init__.py,sha256=k7tL-XgggUwWxHCXyv60rQn-OcXHPg82QcIUkKLEd-c,5057
+ring/__init__.py,sha256=ncBRdHvge6uqpzFyk8_HUQNx4kZxENpVXTJTEK_SjCg,5216
 ring/algebra.py,sha256=F0GwbP8LQP5qGVkoMUYJmkp9Hn2nKAVIkCVYDEjNjGU,3128
 ring/base.py,sha256=Ystn1EjTyOXBhVm5koroV_YPUYtFxrteJLd-XR3kEL8,33840
 ring/maths.py,sha256=qPHH6TpHCK3TgExI98gNEySoSRKOwteN9McUlyUFipI,12207
@@ -6,17 +6,18 @@ 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=bwfVH3qKEnUs6RFgEEeUBnecpBt-nf8cesJbNGDrE7g,28974
+ring/algorithms/jcalc.py,sha256=kQgrRE0XoUBrcdeFbUw_xKnL4m1P2G1Q1m4n7BX2yDk,30064
 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=fzeE7TdUhmGgbbFAyis1tKcyQ4Fo8LigDwD3hUVnH_w,316
+ring/algorithms/custom_joints/__init__.py,sha256=3pQ-Is_HBTQDkzESCNg9VfoP8wvseWmooryG8ERnu_A,366
 ring/algorithms/custom_joints/rr_imp_joint.py,sha256=_YJK0p8_0MHFtr1NuGnNZoxTbwaMQyUjYv7EtsPiU3A,2402
 ring/algorithms/custom_joints/rr_joint.py,sha256=jnRtjtOCALMaq2_0bcu2d7qgfQ6etXpoh43MioRaDmY,1000
+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/batch.py,sha256=9yFxVv11hij-fJXGPxA3zEh1bE2_jrZk0R7kyGaiM5c,2551
-ring/algorithms/generator/finalize_fns.py,sha256=559sGXs06n46p-eme0SE8hn0lXwGT0P2r3-52ElTldo,9861
+ring/algorithms/generator/finalize_fns.py,sha256=EZ5p7fuZu0Zd0rHJzCVg3vy3U9ysny6TfQfolIGPERc,10029
 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
@@ -55,7 +56,7 @@ ring/ml/base.py,sha256=lfwEZLBDglOSRWChUHoH1kezefhttPV9TMEpNIqsMNw,9972
 ring/ml/callbacks.py,sha256=W19QF6_uvaNCjs8ObsjNXD7mv9gFgJBixdRSbB_BynE,13301
 ring/ml/ml_utils.py,sha256=1GXJfeoXbwCbRdYA2np3CbJpSupaw4eyf3quh9y4BO0,6462
 ring/ml/optimizer.py,sha256=fWyF__ezUltrA16SLfOC1jvS3zBh9NJsMYa6-V0frhs,4709
-ring/ml/ringnet.py,sha256=Oud23uKmcvFtwNKdEu2KMMvNAFzJM_yBSRNz2a3CjL4,8670
+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/training_loop.py,sha256=CEokvPQuuk_WCd-J60ZDodJYcPVvyxLfgXDr_DnbzRI,3359
@@ -84,7 +85,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.15.dist-info/METADATA,sha256=zG-f_woph73I5ErczEeJaYXZLNahC-oNXDtzj26f1Po,3821
-imt_ring-1.6.15.dist-info/WHEEL,sha256=GV9aMThwP_4oNCtvEC2ec3qUYutgWeAzklro_0m4WJQ,91
-imt_ring-1.6.15.dist-info/top_level.txt,sha256=EiT790-lAyi8iwTzJArH3f2k77rwhDn00q-4PlmvDQo,5
-imt_ring-1.6.15.dist-info/RECORD,,
+imt_ring-1.6.17.dist-info/METADATA,sha256=j0IKIyc6qAgz9059Z4-b46hi7qTn_ffI7oLw_OrD_Tk,3833
+imt_ring-1.6.17.dist-info/WHEEL,sha256=P9jw-gEje8ByB7_hXoICnHtVCrEwMQh-630tKvQWehc,91
+imt_ring-1.6.17.dist-info/top_level.txt,sha256=EiT790-lAyi8iwTzJArH3f2k77rwhDn00q-4PlmvDQo,5
+imt_ring-1.6.17.dist-info/RECORD,,

{imt_ring-1.6.15.dist-info → imt_ring-1.6.17.dist-info}/WHEEL

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

ring/__init__.py

@@ -121,6 +121,7 @@ _UNIQUE_ID = None
 def setup(
     rr_joint_kwargs: None | dict = dict(),
     rr_imp_joint_kwargs: None | dict = dict(),
+    rsaddle_joint_kwargs: None | dict = dict(),
     suntay_joint_kwargs: None | dict = None,
     train_timing_start: None | float = None,
     unique_id: None | str = None,
@@ -138,6 +139,9 @@ def setup(
     if rr_imp_joint_kwargs is not None:
         custom_joints.register_rr_imp_joint(**rr_imp_joint_kwargs)
 
+    if rsaddle_joint_kwargs is not None:
+        custom_joints.register_rsaddle_joint(**rsaddle_joint_kwargs)
+
     if suntay_joint_kwargs is not None:
         custom_joints.register_suntay(**suntay_joint_kwargs)
 

ring/algorithms/custom_joints/__init__.py

@@ -1,5 +1,6 @@
 from .rr_imp_joint import register_rr_imp_joint
 from .rr_joint import register_rr_joint
+from .rsaddle_joint import register_rsaddle_joint
 from .suntay import ConstantValue_DrawnFnPair
 from .suntay import GP_DrawFnPair
 from .suntay import MLP_DrawnFnPair

ring/algorithms/custom_joints/rsaddle_joint.py

@@ -0,0 +1,40 @@
+import jax.numpy as jnp
+
+import ring
+from ring import maths
+from ring.algorithms.jcalc import _draw_saddle
+from ring.algorithms.jcalc import _p_control_term_rxyz
+from ring.algorithms.jcalc import _qd_from_q_cartesian
+
+
+def register_rsaddle_joint():
+    def _transform(q, params):
+        axes = params["joint_axes"]
+        rot1 = maths.quat_rot_axis(axes[0], q[0])
+        rot2 = maths.quat_rot_axis(axes[1], q[1])
+        rot = maths.quat_mul(rot2, rot1)
+        return ring.Transform.create(rot=rot)
+
+    def _motion_fn_gen(i: int):
+        def _motion_fn(params):
+            axis = params["joint_axes"][i]
+            return ring.base.Motion.create(ang=axis)
+
+        return _motion_fn
+
+    joint_model = ring.JointModel(
+        _transform,
+        motion=[_motion_fn_gen(i) for i in range(2)],
+        rcmg_draw_fn=_draw_saddle,
+        p_control_term=_p_control_term_rxyz,
+        qd_from_q=_qd_from_q_cartesian,
+        init_joint_params=_draw_random_joint_axes,
+    )
+
+    ring.register_new_joint_type("rsaddle", joint_model, 2, overwrite=True)
+
+
+def _draw_random_joint_axes(key):
+    return dict(
+        joint_axes=maths.rotate(jnp.array([1.0, 0, 0]), maths.quat_random(key, (2,)))
+    )

ring/algorithms/generator/finalize_fns.py

@@ -160,6 +160,7 @@ _P_gains = {
     "spherical": jnp.array(3 * [P_rot]),
     "p3d": jnp.array(3 * [P_pos]),
     "saddle": jnp.array([P_rot, P_rot]),
+    "rsaddle": jnp.array([P_rot, P_rot]),
     "frozen": jnp.array([]),
     "suntay": jnp.array([P_rot]),
 }
@@ -182,13 +183,16 @@ class DynamicalSimulation:
 
     @staticmethod
     def assert_test_system(sys: base.System) -> None:
-        "test that system has no zero mass bodies and no joints without damping"
+        "test that system has no zero mass leaf bodies and no joints without damping"
 
         def f(_, __, n, m, d):
-            assert d.size == 0 or m > 0, (
-                "Dynamic simulation is set to `True` which requires masses >= 0, "
-                f"but found body `{n}` with mass={float(m[0])}. This can lead to NaNs."
-            )
+            is_leaf_body = len(sys.children(n)) == 0
+            if is_leaf_body:
+                assert d.size == 0 or m > 0, (
+                    "Dynamic simulation is set to `True` which requires masses >= 0, "
+                    f"but found body `{n}` with mass={float(m[0])}. This can lead to "
+                    "NaNs."
+                )
 
             assert d.size == 0 or all(d > 0.0), (
                 "Dynamic simulation is set to `True` which requires dampings > 0, "

ring/algorithms/jcalc.py

@@ -1,3 +1,4 @@
+from collections import defaultdict
 from dataclasses import asdict
 from dataclasses import dataclass
 from dataclasses import field
@@ -87,11 +88,39 @@ class MotionConfig:
         assert nomotion_config.is_feasible()
         return nomotion_config
 
+    @staticmethod
+    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.
+        """
+        previous_changes = _overwrite_for_joint_type_changes[joint_type]
+        for change in changes:
+            assert change not in previous_changes, f"For jointtype={joint_type} you "
+            f"previously changed the value={change}. You can't change it again, this "
+            "is not supported."
+        previous_changes.update(changes)
+
+        jm = get_joint_model(joint_type)
+
+        def draw_fn(config, *args):
+            return jm.rcmg_draw_fn(replace(config, **changes), *args)
+
+        register_new_joint_type(
+            joint_type,
+            replace(jm, rcmg_draw_fn=draw_fn),
+            base.Q_WIDTHS[joint_type],
+            base.QD_WIDTHS[joint_type],
+            overwrite=True,
+        )
+
     @staticmethod
     def from_register(name: str) -> "MotionConfig":
         return _registered_motion_configs[name]
 
 
+_overwrite_for_joint_type_changes: dict[str, dict] = defaultdict(lambda: dict())
+
+
 _registered_motion_configs = {
     "hinUndHer": MotionConfig(
         t_min=0.3,
@@ -222,7 +251,9 @@ def _is_feasible_config1(c: MotionConfig) -> bool:
     def inside_box_checks(x_min, x_max, x0_min, x0_max) -> bool:
         return (x0_min >= x_min) and (x0_max <= x_max)
 
-    cond2 = inside_box_checks(c.pos_min, c.pos_max, c.pos0_min, c.pos0_max)
+    cond2 = inside_box_checks(
+        _to_float(c.pos_min, 0.0), _to_float(c.pos_max, 0.0), c.pos0_min, c.pos0_max
+    )
 
     return cond1 and cond2
 

ring/ml/ringnet.py

@@ -87,6 +87,7 @@ def make_ring(
     link_output_normalize: bool = True,
     link_output_transform: Optional[Callable] = None,
     layernorm: bool = True,
+    layernorm_trainable: bool = True,
 ) -> SimpleNamespace:
 
     if link_output_normalize:
@@ -104,7 +105,11 @@ def make_ring(
         )
 
         inner_cell = StackedRNNCell(
-            celltype, hidden_state_dim, stack_rnn_cells, layernorm=layernorm
+            celltype,
+            hidden_state_dim,
+            stack_rnn_cells,
+            layernorm=layernorm,
+            layernorm_trainable=layernorm_trainable,
         )
         send_output = hk.nets.MLP([hidden_state_dim, link_output_dim])
         state = hk.get_state(
@@ -143,6 +148,7 @@ class StackedRNNCell(hk.Module):
         hidden_state_dim,
         stacks: int,
         layernorm: bool = False,
+        layernorm_trainable: bool = True,
         name: str | None = None,
     ):
         super().__init__(name)
@@ -150,6 +156,7 @@ class StackedRNNCell(hk.Module):
 
         self.cells = [cell(hidden_state_dim) for _ in range(stacks)]
         self.layernorm = layernorm
+        self.layernorm_trainable = layernorm_trainable
 
     def __call__(self, x, state):
         output = x
@@ -159,7 +166,9 @@ class StackedRNNCell(hk.Module):
             next_state.append(next_state_i)
 
             if self.layernorm:
-                output = hk.LayerNorm(-1, True, True)(output)
+                output = hk.LayerNorm(
+                    -1, self.layernorm_trainable, self.layernorm_trainable
+                )(output)
 
         return output, jnp.stack(next_state)