imt-ring 1.6.47__py3-none-any.whl → 1.7.1__py3-none-any.whl

{imt_ring-1.6.47.dist-info → imt_ring-1.7.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
-Metadata-Version: 2.2
+Metadata-Version: 2.4
 Name: imt-ring
-Version: 1.6.47
+Version: 1.7.1
 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.47.dist-info → imt_ring-1.7.1.dist-info}/RECORD

@@ -1,27 +1,39 @@
-ring/__init__.py,sha256=H1Rd2uXVkux4Z792XyHIkQ8OpDSZBiPqFwyAFDWDU3E,5260
+ring/__init__.py,sha256=y3LuDekHyOCYdzaEDJM5dodClfderAKH-0ufklrwtHY,5266
 ring/algebra.py,sha256=F0GwbP8LQP5qGVkoMUYJmkp9Hn2nKAVIkCVYDEjNjGU,3128
-ring/base.py,sha256=zromjIuMpNBoyiwHa9OCyZvAz7jHjXHZIdRt8fN8PoA,50481
+ring/base.py,sha256=AkG_Gpk7i2j77MzxzjiolJ9WNGcSq_3aOcpu0l6-0e0,50543
 ring/maths.py,sha256=R22SNQutkf9v7Hp9klo0wvJVIyBQz0O8_5oJaDQcFis,12652
 ring/spatial.py,sha256=nmZ-UhRanhyM34bez8uCS4wMwaKqLkuEbgKGP5XNH60,2351
 ring/algorithms/__init__.py,sha256=IiK9EN5Xgs3dB075-A-H-Yad0Z7vzvKIJF2g6X_-C_8,1224
-ring/algorithms/_random.py,sha256=UMyv-VPZLcErrKqs0XB83QJjs8GrmoNsv-zRSxGXvnI,14490
+ring/algorithms/_random.py,sha256=9mSP7M1On_CBz-dEbvE0iDzl_buf125EiYb1y10Ylxc,15805
 ring/algorithms/dynamics.py,sha256=NFOZawjrFoS5RgiWOpG1pQCC8l7RBOEZIi9ok6gvf9U,12268
-ring/algorithms/jcalc.py,sha256=l6BXOmXwrZ_AKKRm4gEHq_k2LSUQ4wd--1qL1qNTcKk,46794
+ring/algorithms/jcalc.py,sha256=zUMKEndSF6zaU7qHjJEU-iEQ_YAFNTNPcOwzdyu33uk,47632
 ring/algorithms/kinematics.py,sha256=IXeTQ-afzeEzLVmQVQ1oTXJxz_lTwvaWlgHeJkhO_8o,7423
 ring/algorithms/sensors.py,sha256=v_TZMyWjffDpPwoyS1fy8X-1i9y1vDf6mk1EmGS2ztc,18251
 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_imp_joint.py,sha256=1sOej-D3q4bSh5KJqcK5fZA3iOHXUlrlT9u-tIa-D6c,2351
 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=TZG307NqdMiXnNY63xEx8AkAjbQBQ4eO6DQ7R4j4D08,16726
 ring/algorithms/generator/__init__.py,sha256=bF-CW3x2x-o6KWESKy-DuxzZPh3UNSjJb_MaAcSHGsQ,277
-ring/algorithms/generator/base.py,sha256=yPH_RIQPU_nlq58HyZ6T3RUm1S5chA3-Ro__-ArYTq0,22669
+ring/algorithms/generator/base.py,sha256=sLIXfFliRUzUKaf84rBQjsExEfmU3XjENrYGD4fm1Q0,23808
 ring/algorithms/generator/batch.py,sha256=xp1X8oYtwI6l2cH4GRu9zw-P8dnh-X1FWTSyixEfgr8,2652
 ring/algorithms/generator/finalize_fns.py,sha256=ty1NaU-Mghx1RL-voivDjS0TWSKNtjTmbdmBnShhn7k,10398
 ring/algorithms/generator/motion_artifacts.py,sha256=2VJbldVDbI3PSyboshIbtYvSAKzBBwGV7cQfYjqvluM,9167
 ring/algorithms/generator/pd_control.py,sha256=dHnhJZx_FqrHD4xFXpQZH-R7rputFkAVGwoBGccZnz4,5767
 ring/algorithms/generator/setup_fns.py,sha256=MFz3czHBeWs1Zk1A8O02CyQpQ-NCyW9PMpbqmKit6es,1455
 ring/algorithms/generator/types.py,sha256=HjNyATFSLfHkXlzdJhvUkiqnhzpXFDDXmWS3LYBlOtU,721
+ring/extras/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+ring/extras/backend.py,sha256=cKSi9sB59texqKzNVASTDczGKLCBL8VVDiP7TNdj41k,1294
+ring/extras/colab.py,sha256=ZLHwP0jNQUsmZJU4l68a5djULPi6T-jYNNHevjIoMn8,1631
+ring/extras/dataloader.py,sha256=dfNPjnxDoKxWGKSImuJ_49CWgBn73vxSEek8COq9nNk,3749
+ring/extras/dataloader_torch.py,sha256=t2DDiB9ZHb_SzFlVbntCGGIybj4F-NoA0PaB4_afjGw,3983
+ring/extras/hdf5.py,sha256=XPIrwogD-d544yy08UJyfLVp1ZKRUtiZukW7RA8VUxQ,5856
+ring/extras/interactive_viewer.py,sha256=-jpoSsDrdzshZw-_MSI96QLvm9foRWTUS_aa_CNFk74,3867
+ring/extras/normalizer.py,sha256=o26stPP6EHasZQxQX0vKqTrhUNZBaJ2O17L6W_gBMN4,1699
+ring/extras/randomize_sys.py,sha256=G_vBIo0OwQkXL2u0djwbaoaeb02C4LQCTNNloOYIU2M,3699
+ring/extras/torch_loss_fn.py,sha256=1LnWTmtxXPxoQFr4QixW12AjpRUfrseSDBmifhu6ErE,2676
+ring/extras/register_gym_envs/__init__.py,sha256=PtPIRBQJ16339xZ9G9VpvqrvcGbQ_Pk_SUz4tQPa9nQ,94
+ring/extras/register_gym_envs/saddle.py,sha256=tA5CyW_akSXyDm0xJ83CtOrUMVElH0f9vZtEDDJQalI,4422
 ring/io/__init__.py,sha256=1gEJdyDCbldbbm8QeZbLmhzSKmaQ-UqTmQgu4DBH2Z4,328
 ring/io/examples.py,sha256=KLf2iCagvRfjs9MCnQsLUlfGBjrQKrD-Qv8U0TtX6Ek,1114
 ring/io/test_examples.py,sha256=htpnSgLG9Fi9_qwSL4F1yLi9sN7ZUrF8dDmiqU3B510,117
@@ -63,8 +75,8 @@ ring/ml/training_loop.py,sha256=yxuUua_4RExq_0GUYm4eUZJsBmtrwDSVL94bWUpYfdo,3586
 ring/ml/params/0x13e3518065c21cd8.pickle,sha256=Zh2k1zK-TNxJl5F7nyTeQ9001qqRE_dfvaq1HWV287A,9355838
 ring/ml/params/0x1d76628065a71e0f.pickle,sha256=YTNVuvfw-nCRD9BH1PZYcR9uCFpNWDhw8Lc50eDn_EE,9351038
 ring/rendering/__init__.py,sha256=Zf7qOdzK3t2hljIrs5P4zFhzHljLSMRyDDZO2YlZk4k,75
-ring/rendering/base_render.py,sha256=Mv9SRLEmuoPVhi46UIjb6xCkKmbWCwIyENGx7nu9REM,9617
-ring/rendering/mujoco_render.py,sha256=HMvZc04I0-lXPBL3hcnBzV2bNiXQAQM7QcHlG_Obmj4,8757
+ring/rendering/base_render.py,sha256=O8Oo9znAgWRE09R7B2yecpwNDJ5veIRoMci144oHwF8,10554
+ring/rendering/mujoco_render.py,sha256=eCmnnzwVZ3BeIo1INswXMZaZ9TDaF1HO50f70spXX2E,9704
 ring/rendering/vispy_render.py,sha256=6Z6S5LNZ7iy9BN1GVb9EDe-Tix5N_SQ1s7ZsfiTSDEA,10261
 ring/rendering/vispy_visuals.py,sha256=ooBZqppnebeL0ANe6V6zUgnNTtDcdkOsa4vZuM4sx-I,7873
 ring/sim2real/__init__.py,sha256=gCLYg8IoMdzUagzhCFcfjZ5GavtIU772L7HR0G5hUtM,251
@@ -73,21 +85,12 @@ ring/sys_composer/__init__.py,sha256=5J_JJJIHfTPcpxh0v4FqiOs81V1REPUd7pgiw2nAN5E
 ring/sys_composer/delete_sys.py,sha256=cIM9KbyLfg7B9121g7yjzuFbjeNu9cil1dPavAYEgzk,3408
 ring/sys_composer/inject_sys.py,sha256=PLuxLbXU7hPtAsqvpsEim9hkoVE26ddrg3OipZNvnhU,3504
 ring/sys_composer/morph_sys.py,sha256=2GpPtS5hT0eZMptdGpt30Hc97OykJNE67lEVRf7sHrc,12700
-ring/utils/__init__.py,sha256=MHHavc8YfjBlmB-zAV42QEQS_ebW7cy0lhWXEVyQU7s,720
-ring/utils/backend.py,sha256=cKSi9sB59texqKzNVASTDczGKLCBL8VVDiP7TNdj41k,1294
+ring/utils/__init__.py,sha256=Q37bjy2wjRGggd77MHlgl_50i2zOuVnPny4yOLiTe-8,567
 ring/utils/batchsize.py,sha256=uCj8LG7elbjEUUzuK29Z3I9T8bxJTcsybY3DdGeqhQs,1786
-ring/utils/colab.py,sha256=ZLHwP0jNQUsmZJU4l68a5djULPi6T-jYNNHevjIoMn8,1631
-ring/utils/dataloader.py,sha256=dfNPjnxDoKxWGKSImuJ_49CWgBn73vxSEek8COq9nNk,3749
-ring/utils/dataloader_torch.py,sha256=t2DDiB9ZHb_SzFlVbntCGGIybj4F-NoA0PaB4_afjGw,3983
-ring/utils/hdf5.py,sha256=XPIrwogD-d544yy08UJyfLVp1ZKRUtiZukW7RA8VUxQ,5856
-ring/utils/normalizer.py,sha256=o26stPP6EHasZQxQX0vKqTrhUNZBaJ2O17L6W_gBMN4,1699
 ring/utils/path.py,sha256=zRPfxYNesvgefkddd26oar6f9433LkMGkhp9dF3rPUs,1926
-ring/utils/randomize_sys.py,sha256=G_vBIo0OwQkXL2u0djwbaoaeb02C4LQCTNNloOYIU2M,3699
 ring/utils/utils.py,sha256=gKwOXLxWraeZfX6EbBcg3hkq30DcXN0mcRUeOSTNiMo,7336
-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.47.dist-info/METADATA,sha256=4acmyig9LGSCfOTMNNOBLrwOaLbCa2EtYXJEfvuMEpc,5888
-imt_ring-1.6.47.dist-info/WHEEL,sha256=jB7zZ3N9hIM9adW7qlTAyycLYW9npaWKLRzaoVcLKcM,91
-imt_ring-1.6.47.dist-info/entry_points.txt,sha256=npNqSOvNiBR0BNa_GL3J66q8Gky3h0G_PHzHzk8oyE0,66
-imt_ring-1.6.47.dist-info/top_level.txt,sha256=EiT790-lAyi8iwTzJArH3f2k77rwhDn00q-4PlmvDQo,5
-imt_ring-1.6.47.dist-info/RECORD,,
+imt_ring-1.7.1.dist-info/METADATA,sha256=cFujMyd1Xpqa8NiYn3iMqnhPHzGvlb_mLBbIPplQA6M,5887
+imt_ring-1.7.1.dist-info/WHEEL,sha256=SmOxYU7pzNKBqASvQJ7DjX3XGUF92lrGhMb3R6_iiqI,91
+imt_ring-1.7.1.dist-info/entry_points.txt,sha256=npNqSOvNiBR0BNa_GL3J66q8Gky3h0G_PHzHzk8oyE0,66
+imt_ring-1.7.1.dist-info/top_level.txt,sha256=EiT790-lAyi8iwTzJArH3f2k77rwhDn00q-4PlmvDQo,5
+imt_ring-1.7.1.dist-info/RECORD,,

{imt_ring-1.6.47.dist-info → imt_ring-1.7.1.dist-info}/WHEEL

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

ring/__init__.py

@@ -35,12 +35,12 @@ def RING(lam: list[int] | None, Ts: float | None, **kwargs) -> ml.AbstractFilter
         >>> import ring
         >>> import numpy as np
         >>>
-        >>> T  : int       = 30        # sequence length     [s]
-        >>> Ts : float     = 0.01      # sampling interval   [s]
-        >>> B  : int       = 1         # batch size
-        >>> lam: list[int] = [0, 1, 2] # parent array
-        >>> N  : int       = len(lam)  # number of bodies
-        >>> T_i: int       = int(T/Ts) # number of timesteps
+        >>> T  : int       = 30         # sequence length     [s]
+        >>> Ts : float     = 0.01       # sampling interval   [s]
+        >>> B  : int       = 1          # batch size
+        >>> lam: list[int] = [-1, 0, 1] # parent array
+        >>> N  : int       = len(lam)   # number of bodies
+        >>> T_i: int       = int(T/Ts)  # number of timesteps
         >>>
         >>> X = np.zeros((B, T_i, N, 9))
         >>> # where X is structured as follows:

ring/algorithms/_random.py

@@ -41,30 +41,48 @@ def random_angle_over_time(
     cdf_bins_min: int = 5,
     cdf_bins_max: Optional[int] = None,
     interpolation_method: str = "cosine",
+    include_standstills_prob: float = 0.0,  # 0.0 means no standstills
+    include_standstills_t_min: float = 0.5,
+    include_standstills_t_max: float = 5.0,
 ) -> jax.Array:
     def body_fn_outer(val):
         i, t, phi, key_t, key_ang, ANG = val
 
-        key_t, consume_t = random.split(key_t)
+        key_t, consume_t, consume_standstill = random.split(key_t, 3)
         key_ang, consume_ang = random.split(key_ang)
         rom_halfsize_float = _to_float(rom_halfsize, t)
         rom_lower = ANG_0 - rom_halfsize_float
         rom_upper = ANG_0 + rom_halfsize_float
-        dt, phi = _resolve_range_of_motion(
-            range_of_motion,
-            range_of_motion_method,
-            rom_lower,
-            rom_upper,
-            _to_float(dang_min, t),
-            _to_float(dang_max, t),
-            _to_float(delta_ang_min, t),
-            _to_float(delta_ang_max, t),
-            t_min,
-            _to_float(t_max, t),
-            phi,
-            consume_t,
-            consume_ang,
-            max_iter,
+
+        is_standstill = jax.random.bernoulli(
+            consume_standstill, include_standstills_prob
+        )
+        dt, phi = jax.lax.cond(
+            is_standstill,
+            lambda: (
+                jax.random.uniform(
+                    consume_t,
+                    minval=include_standstills_t_min,
+                    maxval=include_standstills_t_max,
+                ),
+                phi,
+            ),
+            lambda: _resolve_range_of_motion(
+                range_of_motion,
+                range_of_motion_method,
+                rom_lower,
+                rom_upper,
+                _to_float(dang_min, t),
+                _to_float(dang_max, t),
+                _to_float(delta_ang_min, t),
+                _to_float(delta_ang_max, t),
+                t_min,
+                _to_float(t_max, t),
+                phi,
+                consume_t,
+                consume_ang,
+                max_iter,
+            ),
         )
         t += dt
 
@@ -119,7 +137,8 @@ def random_angle_over_time(
 
 # APPROVED
 def random_position_over_time(
-    key: random.PRNGKey,
+    key_t: random.PRNGKey,
+    key_value: random.PRNGKey,
     POS_0: float,
     pos_min: float | TimeDependentFloat,
     pos_max: float | TimeDependentFloat,
@@ -135,19 +154,14 @@ def random_position_over_time(
     cdf_bins_min: int = 5,
     cdf_bins_max: Optional[int] = None,
     interpolation_method: str = "cosine",
+    include_standstills_prob: float = 0.0,  # 0.0 means no standstills
+    include_standstills_t_min: float = 0.5,
+    include_standstills_t_max: float = 5.0,
 ) -> jax.Array:
     def body_fn_inner(val):
         i, t, t_pre, x, x_pre, key = val
         dt = t - t_pre
 
-        def sample_dx_squared(key):
-            key, consume = random.split(key)
-            dx = (
-                random.uniform(consume) * (2 * dpos_max * t_max**2)
-                - dpos_max * t_max**2
-            )
-            return key, dx
-
         def sample_dx(key):
             key, consume1, consume2 = random.split(key, 3)
             sign = random.choice(consume1, jnp.array([-1.0, 1.0]))
@@ -182,24 +196,43 @@ def random_position_over_time(
         return jnp.logical_not(break_if_true1 | break_if_true2)
 
     def body_fn_outer(val):
-        i, t, t_pre, x, x_pre, key, POS = val
-        key, consume = random.split(key)
-        t += random.uniform(consume, minval=t_min, maxval=_to_float(t_max, t_pre))
+        i, t, t_pre, x, x_pre, key_t, key_value, POS = val
+        key_t, consume_t, consume_standstill = random.split(key_t, 3)
 
-        # that zero resets the max_it count
-        val_inner = (0, t, t_pre, x, x_pre, key)
-        _, t, t_pre, x, x_pre, key = jax.lax.while_loop(
-            cond_fn_inner, body_fn_inner, val_inner
+        is_standstill = jax.random.bernoulli(
+            consume_standstill, include_standstills_prob
+        )
+
+        def is_standstill_branch():
+            dt = random.uniform(
+                consume_t,
+                minval=include_standstills_t_min,
+                maxval=include_standstills_t_max,
+            )
+            t = t_pre + dt
+            return 0, t, t_pre, x, x_pre, key_value
+
+        def no_standstill_branch():
+            dt = random.uniform(consume_t, minval=t_min, maxval=_to_float(t_max, t_pre))
+            t = t_pre + dt
+            # that zero resets the max_it count
+            val_inner = (0, t, t_pre, x, x_pre, key_value)
+            return jax.lax.while_loop(cond_fn_inner, body_fn_inner, val_inner)
+
+        _, t, t_pre, x, x_pre, key_value = jax.lax.cond(
+            is_standstill,
+            is_standstill_branch,
+            no_standstill_branch,
         )
 
         POS_i = jnp.array([[jnp.floor(t / Ts) * Ts, x]])
         POS = jax.lax.dynamic_update_slice_in_dim(POS, POS_i, start_index=i, axis=0)
         t_pre = t
         x_pre = x
-        return i + 1, t, t_pre, x, x_pre, key, POS
+        return i + 1, t, t_pre, x, x_pre, key_t, key_value, POS
 
     def cond_fn_outer(val):
-        i, t, t_pre, x, x_pre, key, POS = val
+        i, t, t_pre, x, x_pre, key_t, key_value, POS = val
         return t <= T
 
     # preallocate POS array
@@ -207,7 +240,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, POS_0, POS_0, key, POS)
+    val_outer = (1, 0.0, 0.0, POS_0, POS_0, key_t, key_value, 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/custom_joints/rr_imp_joint.py

@@ -23,11 +23,10 @@ def register_rr_imp_joint(
         return ring.Transform.create(rot=rot)
 
     def _draw_rr_imp(config, key_t, key_value, dt, N, _):
-        key_t1, key_t2 = jax.random.split(key_t)
         key_value1, key_value2 = jax.random.split(key_value)
-        q_traj_pri = _draw_rxyz(config, key_t1, key_value1, dt, N, _)
+        q_traj_pri = _draw_rxyz(config, key_t, key_value1, dt, N, _)
         q_traj_res = _draw_rxyz(
-            replace(config_res, T=config.T), key_t2, key_value2, dt, N, _
+            replace(config_res, T=config.T), key_t, key_value2, dt, N, _
         )
         # scale to be within bounds
         q_traj_res = q_traj_res * (jnp.deg2rad(ang_max_deg) / jnp.pi)

ring/algorithms/generator/base.py

@@ -1,5 +1,6 @@
 from dataclasses import replace
 from functools import partial
+import json
 import logging
 import random
 from typing import Callable, Optional
@@ -136,6 +137,14 @@ class RCMG:
               affecting joint motion behavior.
         """  # noqa: E501
 
+        # capture all funtion arguments before creating local variables
+        to_json_kwargs = locals()
+        # the purpose is to not capture the RCMG itself since we want to make it
+        # serialisable in the first place
+        to_json_kwargs.pop("self")
+        to_json_kwargs.pop("sys")
+        to_json_kwargs.pop("config")
+
         # add some default values
         randomize_hz_kwargs_defaults = dict(add_dt=True)
         randomize_hz_kwargs_defaults.update(randomize_hz_kwargs)
@@ -186,6 +195,11 @@ class RCMG:
 
         self._disable_tqdm = disable_tqdm
 
+        # store arguments that fully define the RCMG objects for use in `.to_json`
+        self._to_json_sys = sys
+        self._to_json_mconfig = config
+        self._to_json_kwargs = to_json_kwargs
+
     def _compute_repeats(self, sizes: int | list[int]) -> list[int]:
         "how many times the generators are repeated to create a batch of `sizes`"
 
@@ -355,6 +369,21 @@ class RCMG:
 
         return generator
 
+    def serialise_to_dict(self) -> dict:
+        dict_representation = {
+            "system": [_sys.to_str(warn=False) for _sys in self._to_json_sys],
+            "motion_configs": [_config.__dict__ for _config in self._to_json_mconfig],
+            "kwargs": self._to_json_kwargs,
+        }
+        return dict_representation
+
+    def serialise_to_json(self, path_of_json: str) -> None:
+        with open(path_of_json, "w") as file:
+            json.dump(self.serialise_to_dict(), file, indent=4)
+
+    def from_json(self, path_to_json: str) -> "RCMG":
+        raise NotImplementedError
+
 
 def _copy_dicts(f) -> dict:
     def _f(*args, **kwargs):
@@ -526,7 +555,7 @@ def draw_random_q(
     sys: base.System,
     config: jcalc.MotionConfig,
     N: int | None,
-) -> tuple[types.Xy, types.OutputExtras]:
+) -> tuple[jax.random.PRNGKey, jax.Array]:
 
     key_start = key
     # build generalized coordintes vector `q`

ring/algorithms/jcalc.py

@@ -174,6 +174,16 @@ class MotionConfig:
         default_factory=lambda: dict()
     )
 
+    # fields related to simulating standstills (no motion time periods)
+    # these are "Joint Standstills" so the standstills are calculated on
+    # a joint level, for each joint independently
+    # This means that a `standstills_prob` of 20% means that each joint
+    # has at each dt \in [t_min, t_max] drawing process a probability of
+    # 20% that it will just stay at its current joint value
+    include_standstills_prob: float = 0.0  # in %; 0% means no standstills
+    include_standstills_t_min: float = 0.5
+    include_standstills_t_max: float = 5.0
+
     def is_feasible(self) -> bool:
         return _is_feasible_config1(self)
 
@@ -791,12 +801,15 @@ def _draw_rxyz(
         config.cdf_bins_min,
         config.cdf_bins_max,
         config.interpolation_method,
+        config.include_standstills_prob,
+        config.include_standstills_t_min,
+        config.include_standstills_t_max,
     )
 
 
 def _draw_pxyz(
     config: MotionConfig,
-    _: jax.random.PRNGKey,
+    key_t: jax.random.PRNGKey,
     key_value: jax.random.PRNGKey,
     dt: float | jax.Array,
     N: int | None,
@@ -811,6 +824,7 @@ def _draw_pxyz(
     )
     max_iter = 100
     return _random.random_position_over_time(
+        key_t,
         key_value,
         POS_0,
         config.cor_pos_min if cor else config.pos_min,
@@ -827,6 +841,9 @@ def _draw_pxyz(
         config.cdf_bins_min,
         config.cdf_bins_max,
         config.interpolation_method,
+        config.include_standstills_prob,
+        config.include_standstills_t_min,
+        config.include_standstills_t_max,
     )
 
 
@@ -840,7 +857,6 @@ def _draw_spherical(
 ) -> jax.Array:
     # NOTE: We draw 3 euler angles and then build a quaternion.
     # Not ideal, but i am unaware of a better way.
-    @jax.vmap
     def draw_euler_angles(key_t, key_value):
         return _draw_rxyz(
             config,
@@ -853,8 +869,9 @@ def _draw_spherical(
             free_spherical=True,
         )
 
-    triple = lambda key: jax.random.split(key, 3)
-    euler_angles = draw_euler_angles(triple(key_t), triple(key_value)).T
+    euler_angles = jax.vmap(draw_euler_angles, in_axes=(None, 0))(
+        key_t, jax.random.split(key_value, 3)
+    ).T
     q = maths.quat_euler(euler_angles)
     return q
 
@@ -867,7 +884,6 @@ def _draw_saddle(
     N: int | None,
     _: jax.Array,
 ) -> jax.Array:
-    @jax.vmap
     def draw_euler_angles(key_t, key_value):
         return _draw_rxyz(
             config,
@@ -880,14 +896,15 @@ def _draw_saddle(
             free_spherical=False,
         )
 
-    double = lambda key: jax.random.split(key)
-    yz_euler_angles = draw_euler_angles(double(key_t), double(key_value)).T
+    yz_euler_angles = jax.vmap(draw_euler_angles, in_axes=(None, 0))(
+        key_t, jax.random.split(key_value)
+    ).T
     return yz_euler_angles
 
 
 def _draw_p3d_and_cor(
     config: MotionConfig,
-    _: jax.random.PRNGKey,
+    key_t: jax.random.PRNGKey,
     key_value: jax.random.PRNGKey,
     dt: float | jax.Array,
     N: int | None,
@@ -896,7 +913,7 @@ def _draw_p3d_and_cor(
 ) -> jax.Array:
     keys = jax.random.split(key_value, 3)
 
-    def draw(key, xyz: str):
+    def draw(key_value, xyz: str):
         return _draw_pxyz(
             replace(
                 config,
@@ -905,8 +922,8 @@ def _draw_p3d_and_cor(
                 pos0_min=getattr(config, f"pos0_min_p3d_{xyz}"),
                 pos0_max=getattr(config, f"pos0_max_p3d_{xyz}"),
             ),
-            None,
-            key,
+            key_t,
+            key_value,
             dt,
             N,
             None,
@@ -919,26 +936,26 @@ def _draw_p3d_and_cor(
 
 def _draw_p3d(
     config: MotionConfig,
-    _: jax.random.PRNGKey,
+    key_t: jax.random.PRNGKey,
     key_value: jax.random.PRNGKey,
     dt: float | jax.Array,
     N: int | None,
     __: jax.Array,
 ) -> jax.Array:
-    return _draw_p3d_and_cor(config, _, key_value, dt, N, None, cor=False)
+    return _draw_p3d_and_cor(config, key_t, key_value, dt, N, None, cor=False)
 
 
 def _draw_cor(
     config: MotionConfig,
-    _: jax.random.PRNGKey,
+    key_t: jax.random.PRNGKey,
     key_value: jax.random.PRNGKey,
     dt: float | jax.Array,
     N: int | None,
     __: jax.Array,
 ) -> jax.Array:
     key_value1, key_value2 = jax.random.split(key_value)
-    q_free = _draw_free(config, _, key_value1, dt, N, None)
-    q_p3d = _draw_p3d_and_cor(config, _, key_value2, dt, N, None, cor=True)
+    q_free = _draw_free(config, key_t, key_value1, dt, N, None)
+    q_p3d = _draw_p3d_and_cor(config, key_t, key_value2, dt, N, None, cor=True)
     return jnp.concatenate((q_free, q_p3d), axis=1)
 
 
@@ -952,7 +969,7 @@ def _draw_free(
 ) -> jax.Array:
     key_value1, key_value2 = jax.random.split(key_value)
     q = _draw_spherical(config, key_t, key_value1, dt, N, None)
-    pos = _draw_p3d(config, None, key_value2, dt, N, None)
+    pos = _draw_p3d(config, key_t, key_value2, dt, N, None)
     return jnp.concatenate((q, pos), axis=1)
 
 
@@ -975,7 +992,7 @@ def _draw_free_2d(
         enable_range_of_motion=False,
         free_spherical=True,
     )[:, None]
-    pos_yz = _draw_p3d(config, None, key_value2, dt, N, None)[:, :2]
+    pos_yz = _draw_p3d(config, key_t, key_value2, dt, N, None)[:, :2]
     return jnp.concatenate((angle_x, pos_yz), axis=1)
 
 

ring/base.py

@@ -981,6 +981,7 @@ class System(_Base):
 
     def render(
         self,
+        qs: Optional[jax.Array | list[jax.Array]] = None,
         xs: Optional[Transform | list[Transform]] = None,
         camera: Optional[str] = None,
         show_pbar: bool = True,
@@ -1001,7 +1002,7 @@ class System(_Base):
             list[np.ndarray]: Stacked rendered frames. Length == len(xs).
         """
         return ring.rendering.render(
-            self, xs, camera, show_pbar, backend, render_every_nth, **scene_kwargs
+            self, qs, xs, camera, show_pbar, backend, render_every_nth, **scene_kwargs
         )
 
     def render_prediction(

ring/extras/interactive_viewer.py

@@ -0,0 +1,114 @@
+import multiprocessing
+import time
+from typing import Optional
+
+import fire
+import jax.numpy as jnp
+import numpy as np
+
+import ring
+from ring import System
+
+
+class InteractiveViewer:
+    def __init__(self, sys: ring.System, **scene_kwargs):
+        self._mp_dict = multiprocessing.Manager().dict()
+        self._geom_dict = multiprocessing.Manager().dict()
+        self.update_q(np.array(ring.State.create(sys).q))
+        self.process = multiprocessing.Process(
+            target=self._worker,
+            args=(self._mp_dict, self._geom_dict, sys.to_str(), scene_kwargs),
+        )
+        self.process.start()
+
+    def update_q(self, q: np.ndarray):
+        self._mp_dict["q"] = q
+
+    def make_geometry_transparent(self, body_number: int, geom_number: int):
+        geom_name = f"body{body_number}_geom{geom_number}"
+        # the value is not used
+        self._geom_dict[geom_name] = None
+
+    def _worker(self, mp_dict, geom_dict, sys_str, scene_kwargs):
+        from ring.rendering import base_render
+
+        sys = System.from_str(sys_str)
+        while base_render._scene is None or base_render._scene._renderer.is_alive:
+            sys.render(jnp.array(mp_dict["q"]), interactive=True, **scene_kwargs)
+
+            if len(geom_dict) > 0:
+                model = base_render._scene._model
+                processed = []
+                for geom_name in list(geom_dict.keys()):
+                    # Get the geometry ID
+                    geom_id = model.geom(geom_name).id
+                    # Set transparency to 0 (fully transparent)
+                    model.geom_rgba[geom_id, 3] = 0
+                    print(f"Made geom with name={geom_name} transparent (worker)")
+                    processed.append(geom_name)
+
+                for geom_name in processed:
+                    geom_dict.pop(geom_name)
+
+    def __enter__(self):
+        return self
+
+    def close(self):
+        self.process.terminate()
+        self.process.join()
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        self.close()
+
+
+def _fire_main(path_sys_xml: str, path_qs_np: Optional[str] = None, **scene_kwargs):
+    """View motion given by trajectory of minimal coordinates in interactive viewer.
+
+    Args:
+        path_sys_xml (str): Path to xml file defining the system.
+        path_qs_np (str | None, optional): Path to numpy array containing the timeseries of minimal coordinates with
+            shape (T, DOF) where DOF is equal to `sys.q_size()`. Each minimal coordiante is from parent
+            to child. So for example a `spherical` joint that connects the first body to the worldbody
+            has a minimal coordinate of a quaternion that gives from worldbody to first body. The sampling
+            rate of the motion is inferred from the `sys.dt` attribute. If `None` (default), then simply renders the
+            unarticulated pose of the system.
+    """  # noqa: E501
+
+    sys = ring.System.from_xml(path_sys_xml)
+    if path_qs_np is None:
+        qs = np.array(ring.State.create(sys).q)[None]
+    else:
+        qs: np.ndarray = np.load(path_qs_np)
+
+    assert qs.ndim == 2, f"qs.shape = {qs.shape}"
+    T, Q = qs.shape
+    assert Q == sys.q_size(), f"Q={Q} != sys.q_size={sys.q_size()}"
+    dt_target = sys.dt
+
+    with InteractiveViewer(sys, width=640, height=480, **scene_kwargs) as viewer:
+        dt = dt_target
+        last_t = time.time()
+        t = -1
+
+        while True:
+            t = (t + 1) % T
+
+            while dt < dt_target:
+                time.sleep(0.001)
+                dt = time.time() - last_t
+
+            last_t = time.time()
+            viewer.update_q(qs[t])
+            dt = time.time() - last_t
+
+            # process will be stopped if the window is closed
+            if not viewer.process.is_alive():
+                break
+
+
+def main():
+    fire.Fire(_fire_main)
+
+
+if __name__ == "__main__":
+    main()

ring/extras/torch_loss_fn.py

@@ -0,0 +1,93 @@
+"""This module exports a loss function `loss_fn` for training neural networks that
+output quaternions in PyTorch"""
+
+from typing import Sequence
+
+import torch
+
+
+def quat_mul(u: torch.Tensor, v: torch.Tensor) -> torch.Tensor:
+    "Multiplies two quaternions."
+    q = torch.stack(
+        [
+            u[..., 0] * v[..., 0]
+            - u[..., 1] * v[..., 1]
+            - u[..., 2] * v[..., 2]
+            - u[..., 3] * v[..., 3],
+            u[..., 0] * v[..., 1]
+            + u[..., 1] * v[..., 0]
+            + u[..., 2] * v[..., 3]
+            - u[..., 3] * v[..., 2],
+            u[..., 0] * v[..., 2]
+            - u[..., 1] * v[..., 3]
+            + u[..., 2] * v[..., 0]
+            + u[..., 3] * v[..., 1],
+            u[..., 0] * v[..., 3]
+            + u[..., 1] * v[..., 2]
+            - u[..., 2] * v[..., 1]
+            + u[..., 3] * v[..., 0],
+        ],
+        dim=-1,
+    )
+    return q
+
+
+def quat_inv(q: torch.Tensor):
+    return torch.concat([q[..., :1], -q[..., 1:]], dim=-1)
+
+
+def wrap_to_pi(phi):
+    "Wraps angle `phi` (radians) to interval [-pi, pi]."
+    return (phi + torch.pi) % (2 * torch.pi) - torch.pi
+
+
+def quat_angle(q: torch.Tensor):
+    phi = 2 * torch.arctan2(torch.norm(q[..., 1:], dim=-1), q[..., 0])
+    return wrap_to_pi(phi)
+
+
+def safe_normalize(x):
+    return x / (1e-6 + torch.norm(x, dim=-1, keepdim=True))
+
+
+def quat_qrel(q1, q2):
+    "q1^-1 * q2"
+    return quat_mul(quat_inv(q1), q2)
+
+
+@torch.jit.script
+def angle_error(q, qhat):
+    "Absolute angle error in radians"
+    return torch.abs(quat_angle(quat_qrel(q, qhat)))
+
+
+@torch.jit.script
+def inclination_error(q, qhat):
+    "Absolute inclination error in radians. `q`s are from body-to-eps"
+    q_rel = quat_mul(q, quat_inv(qhat))
+    phi_pri = 2 * torch.arctan2(q_rel[..., 3], q_rel[..., 0])
+    q_pri = torch.zeros_like(q)
+    q_pri[..., 0] = torch.cos(phi_pri / 2)
+    q_pri[..., 3] = torch.sin(phi_pri / 2)
+    q_res = quat_mul(q_rel, quat_inv(q_pri))
+    return torch.abs(quat_angle(q_res))
+
+
+def loss_fn(lam: Sequence[int], q: torch.Tensor, qhat: torch.Tensor) -> torch.Tensor:
+    "(..., N, 4) -> (..., N)"
+    *batch_dims, N, F = q.shape
+    assert q.shape == qhat.shape
+    assert F == 4
+    assert N == len(lam)
+    permu = list(reversed(range(q.ndim - 1)))
+    loss_incl = inclination_error(q, qhat).permute(*permu)
+    loss_mae = angle_error(q, qhat).permute(*permu)
+    lam = torch.tensor(lam, device=q.device)
+    return torch.where(
+        lam.reshape(-1, *[1] * len(batch_dims)) == -1, loss_incl, loss_mae
+    ).permute(*permu)
+
+
+def quat_rand(*size: tuple[int]):
+    qs = torch.randn(size=size + (4,))
+    return qs / torch.norm(qs, dim=-1, keepdim=True)

ring/rendering/base_render.py

@@ -1,3 +1,4 @@
+from functools import partial
 from typing import Optional
 
 import jax
@@ -93,15 +94,29 @@ def _load_scene(sys, backend, **scene_kwargs):
     return _scene
 
 
+@jax.jit
+def _jit_forward_kinematics(sys):
+    _, state = kinematics.forward_kinematics(sys, base.State.create(sys))
+    return state.x
+
+
+@jax.jit
+@partial(jax.vmap, in_axes=(None, 0))
+def _jit_vmap_forward_kinematics(sys, q):
+    _, state = kinematics.forward_kinematics(sys, base.State.create(sys, q=q))
+    return state.x
+
+
 def render(
     sys: base.System,
+    qs: Optional[jax.Array | list[jax.Array]] = None,
     xs: Optional[base.Transform | list[base.Transform]] = None,
     camera: Optional[str] = None,
     show_pbar: bool = True,
     backend: str = "mujoco",
     render_every_nth: int = 1,
     **scene_kwargs,
-) -> list[np.ndarray]:
+) -> list[np.ndarray | None]:
     """Render frames from system and trajectory of maximal coordinates `xs`.
 
     Args:
@@ -114,9 +129,18 @@ def render(
     Returns:
         list[np.ndarray]: Stacked rendered frames. Length == len(xs).
     """
+    assert not (qs is not None and xs is not None)
 
+    if xs is None and qs is None:
+        xs = _jit_forward_kinematics(sys)
     if xs is None:
-        xs = kinematics.forward_kinematics(sys, base.State.create(sys))[1].x
+        # throw error if `xs` has been given by accident as `qs` argument
+        qs = utils.to_list(qs)
+        assert not isinstance(
+            qs[0], base.Transform
+        ), "`qs` should be `jax.Array` and not `Transform`; maybe you want to pass `xs` as keyword argument `xs=xs`?"  # noqa: E501
+        qs = jnp.stack(qs, axis=0)
+        xs = _jit_vmap_forward_kinematics(sys, qs)
 
     # convert time-axis of batched xs object into a list of unbatched x objects
     if isinstance(xs, base.Transform) and xs.ndim() == 3:
@@ -144,6 +168,9 @@ def render(
 
     scene = _load_scene(sys, backend, **scene_kwargs)
 
+    if scene_kwargs.get("interactive", False):
+        show_pbar = False
+
     frames = []
     for x in tqdm.tqdm(xs, "Rendering frames..", disable=not show_pbar):
         scene.update(x)
@@ -241,7 +268,7 @@ def render_prediction(
         sys, xs, yhat, transparent_segment_to_root, offset_truth, offset_pred
     )
 
-    frames = render(sys_render, xs_render, **kwargs)
+    frames = render(sys=sys_render, xs=xs_render, **kwargs)
     return frames
 
 

ring/rendering/mujoco_render.py

@@ -118,8 +118,10 @@ def _xml_str_one_body(
     body_number: int, geoms: list[base.Geometry], cameras: list[str], lights: list[str]
 ) -> str:
     inside_body_geoms = ""
-    for geom in geoms:
-        inside_body_geoms += _xml_str_one_geom(geom)
+    for geom_number, geom in enumerate(geoms):
+        inside_body_geoms += _xml_str_one_geom(
+            geom, name=f"body{body_number}_geom{geom_number}"
+        )
 
     inside_body_cameras = ""
     for camera in cameras:
@@ -138,7 +140,7 @@ def _xml_str_one_body(
 """
 
 
-def _xml_str_one_geom(geom: base.Geometry) -> str:
+def _xml_str_one_geom(geom: base.Geometry, name: str) -> str:
     rgba = f'rgba="{_array_to_str(geom.color)}"'
 
     if isinstance(geom, base.Box):
@@ -158,7 +160,8 @@ def _xml_str_one_geom(geom: base.Geometry) -> str:
 
     rot, pos = maths.quat_inv(geom.transform.rot), geom.transform.pos
     rot, pos = f'pos="{_array_to_str(pos)}"', f'quat="{_array_to_str(rot)}"'
-    return f"<geom {type_size} {rgba} {rot} {pos}/>"
+    name = f'name="{name}"'
+    return f"<geom {type_size} {rgba} {rot} {pos} {name}/>"
 
 
 def _array_to_str(arr: Sequence[float]) -> str:
@@ -181,6 +184,8 @@ class MujocoScene:
         floor_z: float = -0.84,
         floor_material: str = "matplane",
         debug: bool = False,
+        interactive: bool = False,
+        interactive_hide_menu: bool = False,
     ) -> None:
         self.debug = debug
         self.height, self.width = height, width
@@ -195,6 +200,8 @@ class MujocoScene:
         self.show_stars = show_stars
         self.show_floor = show_floor
         self.floor_kwargs = dict(z=floor_z, material=floor_material)
+        self.interactive = interactive
+        self.interactive_hide_menu = interactive_hide_menu
 
     def init(self, geoms: list[base.Geometry]):
         self._parent_ids = list(set([geom.link_idx for geom in geoms]))
@@ -208,7 +215,22 @@ class MujocoScene:
             debug=self.debug,
         )
         self._data = mujoco.MjData(self._model)
-        self._renderer = mujoco.Renderer(self._model, self.height, self.width)
+        if self.interactive:
+            import mujoco_viewer
+
+            self._renderer = mujoco_viewer.MujocoViewer(
+                self._model,
+                self._data,
+                width=self.width,
+                height=self.height,
+                hide_menus=self.interactive_hide_menu,
+            )
+
+            if self.interactive_hide_menu:
+                print("Menu can be shown with key `H` for H(elp)")
+
+        else:
+            self._renderer = mujoco.Renderer(self._model, self.height, self.width)
 
     def update(self, x: base.Transform):
         rot, pos = maths.quat_inv(x.rot), x.pos
@@ -234,6 +256,15 @@ class MujocoScene:
 
         mujoco.mj_forward(self._model, self._data)
 
-    def render(self, camera: Optional[str] = None):
-        self._renderer.update_scene(self._data, camera=-1 if camera is None else camera)
+    def render(self, camera: Optional[str] = None) -> np.ndarray | None:
+        if not self.interactive:
+            self._renderer.update_scene(
+                self._data, camera=-1 if camera is None else camera
+            )
         return self._renderer.render()
+
+    def close(self):
+        self._renderer.close()
+
+    def __del__(self):
+        self.close()

ring/utils/__init__.py

@@ -1,11 +1,7 @@
-from . import randomize_sys
 from .batchsize import batchsize_thresholds
 from .batchsize import distribute_batchsize
 from .batchsize import expand_batchsize
 from .batchsize import merge_batchsize
-from .colab import setup_colab_env
-from .normalizer import make_normalizer_from_generator
-from .normalizer import Normalizer
 from .path import parse_path
 from .utils import dict_to_nested
 from .utils import dict_union