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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: imt-ring
-Version: 1.6.16
+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

{imt_ring-1.6.16.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=Gi2fjZojwptEGT0ChVsZ_ktIfj_j9-3EjPKND7tuHxg,29020
+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=_Ao5ASfY53sNRw2hEBlvpvr55xmcA6BqAoQBR04BWEE,9987
+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
@@ -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.16.dist-info/METADATA,sha256=uD0Ot4vDRqXVKrs0Chy9uHk2omhEk8FowBtjfLJGnXg,3833
-imt_ring-1.6.16.dist-info/WHEEL,sha256=P9jw-gEje8ByB7_hXoICnHtVCrEwMQh-630tKvQWehc,91
-imt_ring-1.6.16.dist-info/top_level.txt,sha256=EiT790-lAyi8iwTzJArH3f2k77rwhDn00q-4PlmvDQo,5
-imt_ring-1.6.16.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,,

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

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,