imt-ring 1.5.1__py3-none-any.whl → 1.6.0__py3-none-any.whl

{imt_ring-1.5.1.dist-info → imt_ring-1.6.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: imt-ring
-Version: 1.5.1
+Version: 1.6.0
 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.5.1.dist-info → imt_ring-1.6.0.dist-info}/RECORD

@@ -1,22 +1,22 @@
 ring/__init__.py,sha256=2v6WHlNPucj1XGhDYw-3AlMQGTqH-e4KYK0IaMnBV5s,4760
 ring/algebra.py,sha256=F0GwbP8LQP5qGVkoMUYJmkp9Hn2nKAVIkCVYDEjNjGU,3128
-ring/base.py,sha256=YFPrUWelWswEhq8x8Byv-5pK64mipiGW6x5IlMr4we4,33803
+ring/base.py,sha256=kzBQ54V2xq4KsqRzflyMQ64V-jl8j7eIAsIPIE0gFDk,33127
 ring/maths.py,sha256=jJr_kr78-XDce8B4tXQ2Li-jBntVQhaS8csxglCsj8A,12193
 ring/spatial.py,sha256=nmZ-UhRanhyM34bez8uCS4wMwaKqLkuEbgKGP5XNH60,2351
 ring/algorithms/__init__.py,sha256=IiK9EN5Xgs3dB075-A-H-Yad0Z7vzvKIJF2g6X_-C_8,1224
-ring/algorithms/_random.py,sha256=M9JQSMXSUARWuzlRLP3Wmkuntrk9LZpP30p4_IPgDB4,13805
+ring/algorithms/_random.py,sha256=fc26yEQjSjtf0NluZ41CyeGIRci0ldrRlThueHR9H7U,14007
 ring/algorithms/dynamics.py,sha256=_TwclBXe6vi5C5iJWAIeUIJEIMHQ_1QTmnHvCEpVO0M,10867
-ring/algorithms/jcalc.py,sha256=seis_VQwSvyrZBtmgKwAgSfT_fhXT4Gcyufp0KCUBME,28094
+ring/algorithms/jcalc.py,sha256=bM8VARgqEiVPy7632geKYGk4MZddZfI8XHdW5kXF3HI,28594
 ring/algorithms/kinematics.py,sha256=DOboHI517Vx0pRJUFZtZPmK_qFaiKiQe-37B-M0aC-c,7422
 ring/algorithms/sensors.py,sha256=MICO9Sn0AfoqRx_9KWR3hufsIID-K6SOIg3oPDgsYMU,17869
 ring/algorithms/custom_joints/__init__.py,sha256=fzeE7TdUhmGgbbFAyis1tKcyQ4Fo8LigDwD3hUVnH_w,316
-ring/algorithms/custom_joints/rr_imp_joint.py,sha256=a3JT0w7pB94kZ95eBR8ZO853eSeyjFoiXmhYlaXoHDE,2392
+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/suntay.py,sha256=7-kym1kMDwqYD_2um1roGcBeB8BlTCPe1wljuNGNARA,16676
+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=LOckKDehObDwgOSu_uBhlqkxdztJ0NHTP2mbxwtEcwQ,13335
-ring/algorithms/generator/batch.py,sha256=Hwh5jYZQEmkx73YaXjWd6sZdikmj43spE7DCzGDHXtE,6637
-ring/algorithms/generator/finalize_fns.py,sha256=0fbtwQw89_w0ytQ_aJ877CZGY5fbtb8sbsRO0O8pT34,9081
+ring/algorithms/generator/base.py,sha256=KQSg9uhhR-rC563busVFx4gJrqOx3BXdaChozO9gwTA,14224
+ring/algorithms/generator/batch.py,sha256=ylootnXmj-JyuB_f5OCknHst9wFKO3gkjQbMrFNXY2g,2513
+ring/algorithms/generator/finalize_fns.py,sha256=L_5wIVA7g0P4P2U6EmgcvsoI-YuF3TOaHBwk5_oEaUU,9077
 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
@@ -53,7 +53,7 @@ ring/io/xml/to_xml.py,sha256=fohb-jWMf2cxVdT5dmknsGyrNMseICSbKEz_urbaWbQ,3407
 ring/ml/__init__.py,sha256=8SZTCs9rJ1kzR0Psh7lUzFhIMhKRPIK41mVfxJAGyMo,1471
 ring/ml/base.py,sha256=-3JQ27zMFESNn5zeNer14GJU2yQgiqDcJUaULOeSyp8,9799
 ring/ml/callbacks.py,sha256=W19QF6_uvaNCjs8ObsjNXD7mv9gFgJBixdRSbB_BynE,13301
-ring/ml/ml_utils.py,sha256=hQEmeZoahdJyFrz0NZXYi1Yijl7GvPBdqwzZBzlUIUM,7638
+ring/ml/ml_utils.py,sha256=GooyH5uxA6cJM7ZcWDUfSkSKq6dg7kCIbhkbjJs_rLw,6674
 ring/ml/optimizer.py,sha256=fWyF__ezUltrA16SLfOC1jvS3zBh9NJsMYa6-V0frhs,4709
 ring/ml/ringnet.py,sha256=rgje5AKUKpT8K-vbE9_SgZ3IijR8TJEHnaqxsE57Mhc,8617
 ring/ml/rnno_v1.py,sha256=T4SKG7iypqn2HBQLKhDmJ2Slj2Z5jtUBHvX_6aL8pyM,1103
@@ -62,7 +62,7 @@ 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
 ring/rendering/__init__.py,sha256=Zf7qOdzK3t2hljIrs5P4zFhzHljLSMRyDDZO2YlZk4k,75
-ring/rendering/base_render.py,sha256=s5dF-GVBqjiWkqVuPQMtTLuM7EtA-YrB7RVWFfIaQ1I,8956
+ring/rendering/base_render.py,sha256=Mv9SRLEmuoPVhi46UIjb6xCkKmbWCwIyENGx7nu9REM,9617
 ring/rendering/mujoco_render.py,sha256=uZ-6s6vshsc49N4xvh5KEWQo1f0DveoZqlJ6sIy1QGI,7912
 ring/rendering/vispy_render.py,sha256=QmRyA7Hqk3uS1SKjcncwc4_vd1m4yWryW2X0i4jRvCw,10260
 ring/rendering/vispy_visuals.py,sha256=ooBZqppnebeL0ANe6V6zUgnNTtDcdkOsa4vZuM4sx-I,7873
@@ -72,7 +72,7 @@ ring/sys_composer/__init__.py,sha256=5J_JJJIHfTPcpxh0v4FqiOs81V1REPUd7pgiw2nAN5E
 ring/sys_composer/delete_sys.py,sha256=cIM9KbyLfg7B9121g7yjzuFbjeNu9cil1dPavAYEgzk,3408
 ring/sys_composer/inject_sys.py,sha256=Mj-q-mUjXKwkg-ol6IQAjf9IJfk7pGhez0_WoTKTgm0,3503
 ring/sys_composer/morph_sys.py,sha256=2GpPtS5hT0eZMptdGpt30Hc97OykJNE67lEVRf7sHrc,12700
-ring/utils/__init__.py,sha256=9ZEooVyri0IWXHA5T-L03vP7aWX0zo8qvfNioGnIAkc,696
+ring/utils/__init__.py,sha256=M9bR1-SYtmF9c4mTRIrGuIQws3K2aKUQxbpltIDkgZQ,739
 ring/utils/backend.py,sha256=cKSi9sB59texqKzNVASTDczGKLCBL8VVDiP7TNdj41k,1294
 ring/utils/batchsize.py,sha256=FbOii7MDP4oPZd9GJOKehFatfnb6WZ0b9z349iZYs1A,1786
 ring/utils/colab.py,sha256=ZLHwP0jNQUsmZJU4l68a5djULPi6T-jYNNHevjIoMn8,1631
@@ -80,8 +80,10 @@ ring/utils/hdf5.py,sha256=BzXwVypZmEZeHVgeGZ78YYdi10NEQtnPhdrb8dQAXo0,5856
 ring/utils/normalizer.py,sha256=67L2BU1MRsMT4pD41ta3JJMppLN0ozFmnwrmXDtnqrQ,1698
 ring/utils/path.py,sha256=hAfSlqRi-ew536RnjDDM7IKapdMJc-EvhrR0Y-BCFWc,1265
 ring/utils/randomize_sys.py,sha256=G_vBIo0OwQkXL2u0djwbaoaeb02C4LQCTNNloOYIU2M,3699
-ring/utils/utils.py,sha256=VkB0Gvmlaz2MZdntgjWA0rOpRkvIRpLWRFgIofoY7hs,5441
-imt_ring-1.5.1.dist-info/METADATA,sha256=C4QFyeh8L1nslQdrwJ6tBtL_dqDRGsUl7W8ZExZA2hc,3104
-imt_ring-1.5.1.dist-info/WHEEL,sha256=Z4pYXqR_rTB7OWNDYFOm1qRk0RX6GFP2o8LgvP453Hk,91
-imt_ring-1.5.1.dist-info/top_level.txt,sha256=EiT790-lAyi8iwTzJArH3f2k77rwhDn00q-4PlmvDQo,5
-imt_ring-1.5.1.dist-info/RECORD,,
+ring/utils/utils.py,sha256=k7t-QxMWrNRnjfNB9rSobmLCmhJigE8__gkT-Il0Ee4,6492
+ring/utils/register_gym_envs/__init__.py,sha256=j1qHllOSh8eC24v2d3WjMeFIP-HpixDxTJYJQkriYO0,98
+ring/utils/register_gym_envs/saddle.py,sha256=tA5CyW_akSXyDm0xJ83CtOrUMVElH0f9vZtEDDJQalI,4422
+imt_ring-1.6.0.dist-info/METADATA,sha256=rselknvDNCopDi3O_BrPrDljdaYCxErD7IOZqcUyJ_I,3104
+imt_ring-1.6.0.dist-info/WHEEL,sha256=Z4pYXqR_rTB7OWNDYFOm1qRk0RX6GFP2o8LgvP453Hk,91
+imt_ring-1.6.0.dist-info/top_level.txt,sha256=EiT790-lAyi8iwTzJArH3f2k77rwhDn00q-4PlmvDQo,5
+imt_ring-1.6.0.dist-info/RECORD,,

ring/algorithms/_random.py

@@ -29,7 +29,8 @@ def random_angle_over_time(
     t_min: float,
     t_max: float | TimeDependentFloat,
     T: float,
-    Ts: float,
+    Ts: float | jax.Array,
+    N: Optional[int] = None,
     max_iter: int = 5,
     randomized_interpolation: bool = False,
     range_of_motion: bool = False,
@@ -84,7 +85,10 @@ def random_angle_over_time(
     )
 
     # resample
-    t = jnp.arange(T, step=Ts)
+    if N is None:
+        t = jnp.arange(T, step=Ts)
+    else:
+        t = jnp.arange(N) * Ts
     if randomized_interpolation:
         q = interpolate(cdf_bins_min, cdf_bins_max, method=interpolation_method)(
             t, ANG[:, 0], ANG[:, 1], consume
@@ -117,7 +121,8 @@ def random_position_over_time(
     t_max: float | TimeDependentFloat,
     T: float,
     Ts: float,
-    max_it: int,
+    N: Optional[int] = None,
+    max_it: int = 100,
     randomized_interpolation: bool = False,
     cdf_bins_min: int = 5,
     cdf_bins_max: Optional[int] = None,
@@ -203,7 +208,10 @@ def random_position_over_time(
     )
 
     # resample
-    t = jnp.arange(T, step=Ts)
+    if N is None:
+        t = jnp.arange(T, step=Ts)
+    else:
+        t = jnp.arange(N) * Ts
     if randomized_interpolation:
         r = interpolate(cdf_bins_min, cdf_bins_max, method=interpolation_method)(
             t, POS[:, 0], POS[:, 1], consume

ring/algorithms/custom_joints/rr_imp_joint.py

@@ -2,6 +2,7 @@ from dataclasses import replace
 
 import jax
 import jax.numpy as jnp
+
 import ring
 from ring import maths
 from ring.algorithms.jcalc import _draw_rxyz
@@ -21,12 +22,12 @@ def register_rr_imp_joint(
         rot = ring.maths.quat_mul(rot_res, rot_pri)
         return ring.Transform.create(rot=rot)
 
-    def _draw_rr_imp(config, key_t, key_value, dt, _):
+    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, _)
+        q_traj_pri = _draw_rxyz(config, key_t1, key_value1, dt, N, _)
         q_traj_res = _draw_rxyz(
-            replace(config_res, T=config.T), key_t2, key_value2, dt, _
+            replace(config_res, T=config.T), key_t2, key_value2, dt, N, _
         )
         # scale to be within bounds
         q_traj_res = q_traj_res * (jnp.deg2rad(ang_max_deg) / jnp.pi)

ring/algorithms/custom_joints/suntay.py

@@ -225,7 +225,8 @@ def register_suntay(sconfig: SuntayConfig, name: str = "suntay"):
         mconfig: ring.MotionConfig,
         key_t: jax.random.PRNGKey,
         key_value: jax.random.PRNGKey,
-        dt: float,
+        dt: float | jax.Array,
+        N: int | None,
         _: jax.Array,
     ) -> jax.Array:
         key_value, consume = jax.random.split(key_value)
@@ -251,6 +252,7 @@ def register_suntay(sconfig: SuntayConfig, name: str = "suntay"):
             mconfig.t_max,
             mconfig.T,
             dt,
+            N,
             5,
             mconfig.randomized_interpolation_angle,
             mconfig.range_of_motion_hinge,

ring/algorithms/generator/base.py

@@ -1,3 +1,4 @@
+import random
 from typing import Callable, Optional
 import warnings
 
@@ -33,6 +34,8 @@ class RCMG:
         randomize_positions: bool = False,
         randomize_motion_artifacts: bool = False,
         randomize_joint_params: bool = False,
+        randomize_hz: bool = False,
+        randomize_hz_kwargs: dict = dict(),
         imu_motion_artifacts: bool = False,
         imu_motion_artifacts_kwargs: dict = dict(),
         dynamic_simulation: bool = False,
@@ -68,6 +71,8 @@ class RCMG:
                     randomize_positions=randomize_positions,
                     randomize_motion_artifacts=randomize_motion_artifacts,
                     randomize_joint_params=randomize_joint_params,
+                    randomize_hz=randomize_hz,
+                    randomize_hz_kwargs=randomize_hz_kwargs,
                     imu_motion_artifacts=imu_motion_artifacts,
                     imu_motion_artifacts_kwargs=imu_motion_artifacts_kwargs,
                     dynamic_simulation=dynamic_simulation,
@@ -172,35 +177,37 @@ class RCMG:
         sizes: int | list[int] = 1,
         seed: int = 1,
         shuffle: bool = True,
+        transform=None,
     ) -> types.BatchedGenerator:
         data = self.to_list(sizes, seed)
         assert len(data) >= batchsize
-
-        def data_fn(indices: list[int]):
-            return tree_utils.tree_batch([data[i] for i in indices])
-
-        return batch.generator_from_data_fn(
-            data_fn, list(range(len(data))), shuffle, batchsize
-        )
+        return self.eager_gen_from_list(data, batchsize, shuffle, transform)
 
     @staticmethod
-    def eager_gen_from_paths(
-        paths: str | list[str],
+    def eager_gen_from_list(
+        data: list[tree_utils.PyTree],
         batchsize: int,
-        include_samples: Optional[list[int]] = None,
         shuffle: bool = True,
-        load_all_into_memory: bool = False,
-        tree_transform=None,
-    ) -> tuple[types.BatchedGenerator, int]:
-        paths = utils.to_list(paths)
-        return batch.generator_from_paths(
-            paths,
-            batchsize,
-            include_samples,
-            shuffle,
-            load_all_into_memory=load_all_into_memory,
-            tree_transform=tree_transform,
-        )
+        transform=None,
+    ) -> types.BatchedGenerator:
+        data = data.copy()
+        n_batches, i = len(data) // batchsize, 0
+
+        def generator(key: jax.Array):
+            nonlocal i
+            if shuffle and i == 0:
+                random.shuffle(data)
+
+            start, stop = i * batchsize, (i + 1) * batchsize
+            batch = tree_utils.tree_batch(data[start:stop], backend="numpy")
+            batch = utils.pytree_deepcopy(batch)
+            if transform is not None:
+                batch = transform(batch)
+
+            i = (i + 1) % n_batches
+            return batch
+
+        return generator
 
 
 def _copy_dicts(f) -> dict:
@@ -229,6 +236,8 @@ def _build_mconfig_batched_generator(
     randomize_positions: bool,
     randomize_motion_artifacts: bool,
     randomize_joint_params: bool,
+    randomize_hz: bool,
+    randomize_hz_kwargs: dict,
     imu_motion_artifacts: bool,
     imu_motion_artifacts_kwargs: dict,
     dynamic_simulation: bool,
@@ -318,16 +327,29 @@ def _build_mconfig_batched_generator(
         key, *consume = jax.random.split(key, len(config) + 1)
         syss = jax.vmap(_setup_fn, (0, None))(jnp.array(consume), sys)
 
+        if randomize_hz:
+            assert "sampling_rates" in randomize_hz_kwargs
+            hzs = randomize_hz_kwargs["sampling_rates"]
+            assert len(set([c.T for c in config])) == 1
+            N = int(min(hzs) * config[0].T)
+            key, consume = jax.random.split(key)
+            dt = 1 / jax.random.choice(consume, jnp.array(hzs))
+            # makes sys.dt from float to AbstractArray
+            syss = syss.replace(dt=jnp.array(dt))
+        else:
+            N = None
+
         qs = []
         for i, _config in enumerate(config):
-            key, _q = draw_random_q(key, syss[i], _config)
+            key, _q = draw_random_q(key, syss[i], _config, N)
             qs.append(_q)
         qs = jnp.stack(qs)
 
         @jax.vmap
         def _vmapped_context(key, q, sys):
             x, _ = jax.vmap(kinematics.forward_kinematics_transforms, (None, 0))(sys, q)
-            Xy, extras = ({}, {}), (key, q, x, sys)
+            X = {"dt": jnp.array(sys.dt)} if randomize_hz else {}
+            Xy, extras = (X, {}), (key, q, x, sys)
             return _finalize_fn(Xy, extras)
 
         keys = jax.random.split(key, len(config))
@@ -343,6 +365,7 @@ def draw_random_q(
     key: types.PRNGKey,
     sys: base.System,
     config: jcalc.MotionConfig,
+    N: int | None,
 ) -> tuple[types.Xy, types.OutputExtras]:
 
     key_start = key
@@ -363,7 +386,7 @@ 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, joint_params)
+        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/batch.py

@@ -1,7 +1,3 @@
-from pathlib import Path
-import random
-from typing import Optional
-
 import jax
 import jax.numpy as jnp
 import numpy as np
@@ -88,142 +84,3 @@ def generators_eager_to_list(
             data.extend([jax.tree_map(lambda a: a[i], sample) for i in range(size)])
 
     return data
-
-
-def _is_nan(ele: tree_utils.PyTree, i: int, verbose: bool = False):
-    isnan = np.any([np.any(np.isnan(arr)) for arr in jax.tree_util.tree_leaves(ele)])
-    if isnan:
-        X, y = ele
-        dt = X["dt"].flatten()[0]
-        if verbose:
-            print(f"Sample with idx={i} is nan. It will be replaced. (dt={dt})")
-        return True
-    return False
-
-
-def _replace_elements_w_nans(list_of_data: list, include_samples: list[int]) -> list:
-    list_of_data_nonan = []
-    for i, ele in enumerate(list_of_data):
-        if _is_nan(ele, i, verbose=True):
-            while True:
-                j = random.choice(include_samples)
-                if not _is_nan(list_of_data[j], j):
-                    ele = list_of_data[j]
-                    break
-        list_of_data_nonan.append(ele)
-    return list_of_data_nonan
-
-
-_list_of_data = None
-_paths = None
-
-
-def _data_fn_from_paths(
-    paths: list[str],
-    include_samples: list[int] | None,
-    load_all_into_memory: bool,
-    tree_transform,
-):
-    "`data_fn` returns numpy arrays."
-    global _list_of_data, _paths
-
-    # expanduser
-    paths = [utils.parse_path(p, mkdir=False) for p in paths]
-    extensions = list(set([Path(p).suffix for p in paths]))
-    assert len(extensions) == 1, f"{extensions}"
-    h5 = extensions[0] == ".h5"
-
-    if h5 and not load_all_into_memory:
-
-        def data_fn(indices: list[int]):
-            tree = utils.hdf5_load_from_multiple(paths, indices)
-            return tree if tree_transform is None else tree_transform(tree)
-
-        N = sum([utils.hdf5_load_length(p) for p in paths])
-    else:
-
-        load_from_path = utils.hdf5_load if h5 else utils.pickle_load
-
-        def load_fn(path):
-            tree = load_from_path(path)
-            tree = tree if tree_transform is None else tree_transform(tree)
-            return [
-                jax.tree_map(lambda arr: arr[i], tree)
-                for i in range(tree_utils.tree_shape(tree))
-            ]
-
-        if paths != _paths or len(_list_of_data) == 0:
-            _paths = paths
-
-            _list_of_data = []
-            for p in paths:
-                _list_of_data += load_fn(p)
-
-        N = len(_list_of_data)
-        list_of_data = _replace_elements_w_nans(
-            _list_of_data,
-            include_samples if include_samples is not None else list(range(N)),
-        )
-
-        if include_samples is not None:
-            list_of_data = [
-                ele if i in include_samples else None
-                for i, ele in enumerate(list_of_data)
-            ]
-
-        def data_fn(indices: list[int]):
-            return tree_utils.tree_batch(
-                [list_of_data[i] for i in indices], backend="numpy"
-            )
-
-    if include_samples is None:
-        include_samples = list(range(N))
-
-    return data_fn, include_samples.copy()
-
-
-def generator_from_data_fn(
-    data_fn,
-    include_samples: list[int],
-    shuffle: bool,
-    batchsize: int,
-) -> types.BatchedGenerator:
-    # such that we don't mutate out of scope
-    include_samples = include_samples.copy()
-
-    N = len(include_samples)
-    n_batches, i = N // batchsize, 0
-
-    def generator(key: jax.Array):
-        nonlocal i
-        if shuffle and i == 0:
-            random.shuffle(include_samples)
-
-        start, stop = i * batchsize, (i + 1) * batchsize
-        batch = data_fn(include_samples[start:stop])
-
-        i = (i + 1) % n_batches
-        return utils.pytree_deepcopy(batch)
-
-    return generator
-
-
-def generator_from_paths(
-    paths: list[str],
-    batchsize: int,
-    include_samples: Optional[list[int]] = None,
-    shuffle: bool = True,
-    load_all_into_memory: bool = False,
-    tree_transform=None,
-) -> tuple[types.BatchedGenerator, int]:
-    "Returns: gen, where gen(key) -> Pytree[numpy]"
-    data_fn, include_samples = _data_fn_from_paths(
-        paths, include_samples, load_all_into_memory, tree_transform
-    )
-
-    N = len(include_samples)
-    assert N >= batchsize
-
-    generator = generator_from_data_fn(data_fn, include_samples, shuffle, batchsize)
-
-    return generator, N

ring/algorithms/generator/finalize_fns.py

@@ -251,8 +251,8 @@ def _expand_dt(X: dict, T: int):
     return X
 
 
-def _expand_then_flatten(args):
-    X, y = args
+def _expand_then_flatten(Xy):
+    X, y = Xy
     gyr = X["0"]["gyr"]
 
     batched = True

ring/algorithms/jcalc.py

@@ -274,8 +274,15 @@ def join_motionconfigs(
 
 
 DRAW_FN = Callable[
-    # config, key_t, key_value, dt, params
-    [MotionConfig, jax.random.PRNGKey, jax.random.PRNGKey, float, jax.Array],
+    # config, key_t, key_value, dt, N, params
+    [
+        MotionConfig,
+        jax.random.PRNGKey,
+        jax.random.PRNGKey,
+        float | jax.Array,
+        int | None,
+        jax.Array,
+    ],
     jax.Array,
 ]
 P_CONTROL_TERM = Callable[
@@ -410,7 +417,8 @@ def _draw_rxyz(
     config: MotionConfig,
     key_t: jax.random.PRNGKey,
     key_value: jax.random.PRNGKey,
-    dt: float,
+    dt: float | jax.Array,
+    N: int | None,
     _: jax.Array,
     # TODO, delete these args and pass a modifified `config` with `replace` instead
     enable_range_of_motion: bool = True,
@@ -435,6 +443,7 @@ def _draw_rxyz(
         config.t_max,
         config.T,
         dt,
+        N,
         max_iter,
         config.randomized_interpolation_angle,
         config.range_of_motion_hinge if enable_range_of_motion else False,
@@ -449,7 +458,8 @@ def _draw_pxyz(
     config: MotionConfig,
     _: jax.random.PRNGKey,
     key_value: jax.random.PRNGKey,
-    dt: float,
+    dt: float | jax.Array,
+    N: int | None,
     __: jax.Array,
     cor: bool = False,
 ) -> jax.Array:
@@ -467,6 +477,7 @@ def _draw_pxyz(
         config.cor_t_max if cor else config.t_max,
         config.T,
         dt,
+        N,
         max_iter,
         config.randomized_interpolation_position,
         config.cdf_bins_min,
@@ -479,7 +490,8 @@ def _draw_spherical(
     config: MotionConfig,
     key_t: jax.random.PRNGKey,
     key_value: jax.random.PRNGKey,
-    dt: float,
+    dt: float | jax.Array,
+    N: int | None,
     _: jax.Array,
 ) -> jax.Array:
     # NOTE: We draw 3 euler angles and then build a quaternion.
@@ -491,6 +503,7 @@ def _draw_spherical(
             key_t,
             key_value,
             dt,
+            N,
             None,
             enable_range_of_motion=False,
             free_spherical=True,
@@ -506,7 +519,8 @@ def _draw_saddle(
     config: MotionConfig,
     key_t: jax.random.PRNGKey,
     key_value: jax.random.PRNGKey,
-    dt: float,
+    dt: float | jax.Array,
+    N: int | None,
     _: jax.Array,
 ) -> jax.Array:
     @jax.vmap
@@ -516,6 +530,7 @@ def _draw_saddle(
             key_t,
             key_value,
             dt,
+            N,
             None,
             enable_range_of_motion=False,
             free_spherical=False,
@@ -530,11 +545,12 @@ def _draw_p3d_and_cor(
     config: MotionConfig,
     _: jax.random.PRNGKey,
     key_value: jax.random.PRNGKey,
-    dt: float,
+    dt: float | jax.Array,
+    N: int | None,
     __: jax.Array,
     cor: bool,
 ) -> jax.Array:
-    pos = jax.vmap(lambda key: _draw_pxyz(config, None, key, dt, None, cor))(
+    pos = jax.vmap(lambda key: _draw_pxyz(config, None, key, dt, N, None, cor))(
         jax.random.split(key_value, 3)
     )
     return pos.T
@@ -544,22 +560,24 @@ def _draw_p3d(
     config: MotionConfig,
     _: jax.random.PRNGKey,
     key_value: jax.random.PRNGKey,
-    dt: float,
+    dt: float | jax.Array,
+    N: int | None,
     __: jax.Array,
 ) -> jax.Array:
-    return _draw_p3d_and_cor(config, _, key_value, dt, None, cor=False)
+    return _draw_p3d_and_cor(config, _, key_value, dt, N, None, cor=False)
 
 
 def _draw_cor(
     config: MotionConfig,
     _: jax.random.PRNGKey,
     key_value: jax.random.PRNGKey,
-    dt: float,
+    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, None)
-    q_p3d = _draw_p3d_and_cor(config, _, key_value2, dt, None, cor=True)
+    q_free = _draw_free(config, _, key_value1, dt, N, None)
+    q_p3d = _draw_p3d_and_cor(config, _, key_value2, dt, N, None, cor=True)
     return jnp.concatenate((q_free, q_p3d), axis=1)
 
 
@@ -567,12 +585,13 @@ def _draw_free(
     config: MotionConfig,
     key_t: jax.random.PRNGKey,
     key_value: jax.random.PRNGKey,
-    dt: float,
+    dt: float | jax.Array,
+    N: int | None,
     __: jax.Array,
 ) -> jax.Array:
     key_value1, key_value2 = jax.random.split(key_value)
-    q = _draw_spherical(config, key_t, key_value1, dt, None)
-    pos = _draw_p3d(config, None, key_value2, dt, None)
+    q = _draw_spherical(config, key_t, key_value1, dt, N, None)
+    pos = _draw_p3d(config, None, key_value2, dt, N, None)
     return jnp.concatenate((q, pos), axis=1)
 
 
@@ -580,7 +599,8 @@ def _draw_free_2d(
     config: MotionConfig,
     key_t: jax.random.PRNGKey,
     key_value: jax.random.PRNGKey,
-    dt: float,
+    dt: float | jax.Array,
+    N: int | None,
     __: jax.Array,
 ) -> jax.Array:
     key_value1, key_value2 = jax.random.split(key_value)
@@ -589,16 +609,20 @@ def _draw_free_2d(
         key_t,
         key_value1,
         dt,
+        N,
         None,
         enable_range_of_motion=False,
         free_spherical=True,
     )[:, None]
-    pos_yz = _draw_p3d(config, None, key_value2, dt, None)[:, :2]
+    pos_yz = _draw_p3d(config, None, key_value2, dt, N, None)[:, :2]
     return jnp.concatenate((angle_x, pos_yz), axis=1)
 
 
-def _draw_frozen(config: MotionConfig, _, __, dt: float, ___) -> jax.Array:
-    N = int(config.T / dt)
+def _draw_frozen(
+    config: MotionConfig, _, __, dt: float | jax.Array, N: int | None, ___
+) -> jax.Array:
+    if N is None:
+        N = int(config.T / dt)
     return jnp.zeros((N, 0))
 
 

ring/base.py

@@ -490,24 +490,6 @@ class System(_Base):
         new_link_names = [prefix + name + suffix for name in self.link_names]
         return self.replace(link_names=new_link_names)
 
-    @staticmethod
-    def deep_equal(a, b):
-        if type(a) is not type(b):
-            return False
-        if isinstance(a, _Base):
-            return System.deep_equal(a.__dict__, b.__dict__)
-        if isinstance(a, dict):
-            if a.keys() != b.keys():
-                return False
-            return all(System.deep_equal(a[k], b[k]) for k in a.keys())
-        if isinstance(a, (list, tuple)):
-            if len(a) != len(b):
-                return False
-            return all(System.deep_equal(a[i], b[i]) for i in range(len(a)))
-        if isinstance(a, (np.ndarray, jnp.ndarray, jax.Array)):
-            return jnp.array_equal(a, b)
-        return a == b
-
     def _replace_free_with_cor(self) -> "System":
         # check that
         # - all free joints connect to -1

ring/ml/ml_utils.py

@@ -3,17 +3,16 @@ from functools import partial
 import os
 from pathlib import Path
 import pickle
-import random
 import time
 from typing import Optional, Protocol
 import warnings
 
 import jax
 import numpy as np
-import ring
-from ring.utils import import_lib
 from tree_utils import PyTree
 
+import ring
+from ring.utils import import_lib
 import wandb
 
 # An arbitrarily nested dictionary with Array leaves; Or strings
@@ -231,42 +230,5 @@ def save_model_tf(jax_func, path: str, *input, validate: bool = True):
         )
 
 
-def train_val_split(
-    tps: list[str],
-    bs: int,
-    n_batches_for_val: int = 1,
-    transform_gen=None,
-    tree_transform=None,
-):
-    "Uses `random` module for shuffeling."
-    if transform_gen is None:
-        transform_gen = lambda gen: gen
-
-    len_val = n_batches_for_val * bs
-
-    _, N = ring.RCMG.eager_gen_from_paths(tps, 1)
-    include_samples = list(range(N))
-    random.shuffle(include_samples)
-
-    train_data, val_data = include_samples[:-len_val], include_samples[-len_val:]
-    X_val, y_val = transform_gen(
-        ring.RCMG.eager_gen_from_paths(
-            tps, len_val, val_data, tree_transform=tree_transform
-        )[0]
-    )(jax.random.PRNGKey(420))
-
-    generator = transform_gen(
-        ring.RCMG.eager_gen_from_paths(
-            tps,
-            bs,
-            train_data,
-            load_all_into_memory=True,
-            tree_transform=tree_transform,
-        )[0]
-    )
-
-    return generator, (X_val, y_val)
-
-
 def _unknown_link_names(N: int):
     return [f"link{i}" for i in range(N)]

ring/rendering/base_render.py

@@ -44,27 +44,19 @@ _rgbas = {
 }
 
 
-def render(
-    sys: base.System,
-    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]:
-    """Render frames from system and trajectory of maximal coordinates `xs`.
+_args = None
+_scene = None
 
-    Args:
-        sys (base.System): System to render.
-        xs (base.Transform | list[base.Transform]): Single or time-series
-        of maximal coordinates `xs`.
-        show_pbar (bool, optional): Whether or not to show a progress bar.
-        Defaults to True.
 
-    Returns:
-        list[np.ndarray]: Stacked rendered frames. Length == len(xs).
-    """
+def _load_scene(sys, backend, **scene_kwargs):
+    global _args, _scene
+
+    args = (sys, backend, scene_kwargs)
+    if _args is not None:
+        if utils.tree_equal(_args, args):
+            return _scene
+
+    _args = args
     if backend == "mujoco":
         utils.import_lib("mujoco")
         from ring.rendering.mujoco_render import MujocoScene
@@ -95,6 +87,34 @@ def render(
     # convert all colors to rgbas
     geoms_rgba = [_color_to_rgba(geom) for geom in geoms]
 
+    scene.init(geoms_rgba)
+
+    _scene = scene
+    return _scene
+
+
+def render(
+    sys: base.System,
+    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]:
+    """Render frames from system and trajectory of maximal coordinates `xs`.
+
+    Args:
+        sys (base.System): System to render.
+        xs (base.Transform | list[base.Transform]): Single or time-series
+        of maximal coordinates `xs`.
+        show_pbar (bool, optional): Whether or not to show a progress bar.
+        Defaults to True.
+
+    Returns:
+        list[np.ndarray]: Stacked rendered frames. Length == len(xs).
+    """
+
     if xs is None:
         xs = kinematics.forward_kinematics(sys, base.State.create(sys))[1].x
 
@@ -122,7 +142,7 @@ def render(
     for x in xs:
         data_check(x)
 
-    scene.init(geoms_rgba)
+    scene = _load_scene(sys, backend, **scene_kwargs)
 
     frames = []
     for x in tqdm.tqdm(xs, "Rendering frames..", disable=not show_pbar):
@@ -132,19 +152,9 @@ def render(
     return frames
 
 
-def render_prediction(
-    sys: base.System,
-    xs: base.Transform | list[base.Transform],
-    yhat: dict | jax.Array | np.ndarray,
-    # by default we don't predict the global rotation
-    transparent_segment_to_root: bool = True,
-    **kwargs,
+def _render_prediction_internals(
+    sys, xs, yhat, transparent_segment_to_root, offset_truth, offset_pred
 ):
-    "`xs` matches `sys`. `yhat` matches `sys_noimu`. `yhat` are child-to-parent."
-
-    offset_truth = kwargs.pop("offset_truth", [0, 0, 0])
-    offset_pred = kwargs.pop("offset_pred", [0, 0, 0])
-
     if isinstance(xs, list):
         # list -> batched Transform
         xs = xs[0].batch(*xs[1:])
@@ -185,7 +195,7 @@ def render_prediction(
     xs, xshat = xs.transpose((1, 0, 2)), xshat.transpose((1, 0, 2))
 
     add_offset = lambda x, offset: algebra.transform_mul(
-        x, base.Transform.create(pos=jnp.array(offset, dtype=jnp.float32))
+        x, base.Transform.create(pos=offset)
     )
 
     # create mapping from `name` -> Transform
@@ -211,6 +221,26 @@ def render_prediction(
     xs_render = xs_render[0].batch(*xs_render[1:])
     xs_render = xs_render.transpose((1, 0, 2))
 
+    return sys_render, xs_render
+
+
+def render_prediction(
+    sys: base.System,
+    xs: base.Transform | list[base.Transform],
+    yhat: dict | jax.Array | np.ndarray,
+    # by default we don't predict the global rotation
+    transparent_segment_to_root: bool = True,
+    **kwargs,
+):
+    "`xs` matches `sys`. `yhat` matches `sys_noimu`. `yhat` are child-to-parent."
+
+    offset_truth = jnp.array(kwargs.pop("offset_truth", [0.0, 0, 0]))
+    offset_pred = jnp.array(kwargs.pop("offset_pred", [0.0, 0, 0]))
+
+    sys_render, xs_render = jax.jit(_render_prediction_internals, static_argnums=3)(
+        sys, xs, yhat, transparent_segment_to_root, offset_truth, offset_pred
+    )
+
     frames = render(sys_render, xs_render, **kwargs)
     return frames
 

ring/utils/__init__.py

@@ -16,6 +16,7 @@ from .utils import pickle_load
 from .utils import pickle_save
 from .utils import primes
 from .utils import pytree_deepcopy
+from .utils import replace_elements_w_nans
 from .utils import sys_compare
 from .utils import to_list
 from .utils import tree_equal

ring/utils/register_gym_envs/__init__.py

@@ -0,0 +1,3 @@
+import gymnasium as gym
+
+gym.register("Saddle-v0", "src.ring.utils.register_gym_envs.saddle:Env")

ring/utils/register_gym_envs/saddle.py

@@ -0,0 +1,109 @@
+from gymnasium import spaces
+import gymnasium as gym
+import jax
+import numpy as np
+
+import ring
+
+xml = """
+<x_xy model="lam2">
+  <options dt="0.01" gravity="0.0 0.0 9.81"/>
+  <worldbody>
+    <body joint="free" name="seg1" pos="0.4 0.0 0.0" pos_min="0.2 -0.05 -0.05" pos_max="0.55 0.05 0.05" damping="5.0 5.0 5.0 25.0 25.0 25.0">
+      <geom pos="0.1 0.0 0.0" mass="1.0" color="dustin_exp_blue" edge_color="black" type="box" dim="0.2 0.05 0.05"/>
+      <geom pos="0.05 0.05 0.0" mass="0.1" color="black" edge_color="black" type="box" dim="0.01 0.1 0.01"/>
+      <geom pos="0.15 -0.05 0.0" mass="0.1" color="black" edge_color="black" type="box" dim="0.01 0.1 0.01"/>
+      <body joint="frozen" name="imu1" pos="0.099999994 0.0 0.035" pos_min="0.050000012 -0.05 -0.05" pos_max="0.15 0.05 0.05">
+        <geom mass="0.1" color="dustin_exp_orange" edge_color="black" type="box" dim="0.05 0.03 0.02"/>
+      </body>
+      <body joint="saddle" name="seg2" pos="0.20000002 0.0 0.0" pos_min="0.0 -0.05 -0.05" pos_max="0.35 0.05 0.05" damping="3.0 3.0">
+        <geom pos="0.1 0.0 0.0" mass="1.0" color="dustin_exp_blue" edge_color="black" type="box" dim="0.2 0.05 0.05"/>
+        <geom pos="0.1 0.05 0.0" mass="0.1" color="black" edge_color="black" type="box" dim="0.01 0.1 0.01"/>
+        <geom pos="0.15 -0.05 0.0" mass="0.1" color="black" edge_color="black" type="box" dim="0.01 0.1 0.01"/>
+        <body joint="frozen" name="imu2" pos="0.100000024 0.0 0.035" pos_min="0.050000012 -0.05 -0.05" pos_max="0.14999998 0.05 0.05">
+          <geom mass="0.1" color="dustin_exp_orange" edge_color="black" type="box" dim="0.05 0.03 0.02"/>
+        </body>
+      </body>
+    </body>
+  </worldbody>
+</x_xy>
+"""  # noqa: E501
+
+
+class Env(gym.Env):
+    metadata = {"render_modes": ["rgb_array"], "render_fps": 25}
+
+    def __init__(self, T: float = 60):
+        self._sys = ring.System.create(xml)
+        self._generator = ring.RCMG(
+            self._sys,
+            ring.MotionConfig(T=T, pos_min=0),
+            add_X_imus=1,
+            # child-to-parent
+            add_y_relpose=1,
+            cor=True,
+            disable_tqdm=True,
+            keep_output_extras=True,
+        ).to_lazy_gen()
+        # warmup jit compile
+        self._generator(jax.random.PRNGKey(1))
+
+        self.observation_space = spaces.Box(-float("inf"), float("inf"), shape=(12,))
+        # quaternion; from seg2 to seg1, so child-to-parent
+        self.action_space = spaces.Box(-1.0, 1.0, shape=(4,))
+        self.reward_range = (-float("inf"), 0.0)
+
+        self._action = None
+
+    def reset(self, seed=None, options=None):
+        super().reset(seed=seed, options=options)
+
+        jax_seed = self.np_random.integers(1, int(1e18))
+        (X, y), (_, _, xs, _) = self._generator(jax.random.PRNGKey(jax_seed))
+        self._xs = xs[0]
+        self._truth = y["seg2"][0]
+        self._T = self._truth.shape[0]
+        self._observations = np.zeros((self._T, 12), dtype=np.float32)
+        self._observations[:, :3] = X["seg1"]["acc"][0]
+        self._observations[:, 3:6] = X["seg1"]["gyr"][0]
+        self._observations[:, 6:9] = X["seg2"]["acc"][0]
+        self._observations[:, 9:12] = X["seg2"]["gyr"][0]
+        self._t = 0
+
+        return self._get_obs(), self._get_info()
+
+    def _get_obs(self):
+        return self._observations[self._t]
+
+    def _get_info(self):
+        return {"truth": self._truth[self._t]}
+
+    def step(self, action):
+        self._t += 1
+
+        # convert to unit quaternion
+        self._action = action / np.linalg.norm(action)
+        reward = -self._abs_angle(self._truth[self._t - 1], self._action)
+
+        terminated = False
+        truncated = self._t >= (self._T - 1)
+
+        return self._get_obs(), reward, terminated, truncated, self._get_info()
+
+    def _abs_angle(self, q, qhat) -> float:
+        return float(jax.jit(ring.maths.angle_error)(q, qhat))
+
+    def render(self):
+        light = '<light pos="0 0 3" dir="0 0 -1" directional="false"/>'
+        render_kwargs = dict(
+            show_pbar=False,
+            camera="target",
+            width=640,
+            height=480,
+            add_lights={-1: light},
+        )
+        x = [self._xs[self._t]]
+        if self._action is None:
+            return self._sys.render(x, **render_kwargs)[0]
+        yhat = {"seg1": np.array([[1.0, 0, 0, 0]]), "seg2": self._action[None]}
+        return self._sys.render_prediction(x, yhat, **render_kwargs)[0]

ring/utils/utils.py

@@ -1,11 +1,13 @@
 from importlib import import_module as _import_module
 import io
 import pickle
+import random
 from typing import Optional
 
 import jax
 import jax.numpy as jnp
 import numpy as np
+import tree_utils
 
 from ring.base import _Base
 from ring.base import Geometry
@@ -14,7 +16,6 @@ from .path import parse_path
 
 
 def tree_equal(a, b):
-    "Copied from Marcel / Thomas"
     if type(a) is not type(b):
         return False
     if isinstance(a, _Base):
@@ -181,3 +182,36 @@ def gcd(a: int, b: int) -> int:
     while b:
         a, b = b, a % b
     return a
+
+
+def replace_elements_w_nans(
+    list_of_data: list[tree_utils.PyTree],
+    include_elements: Optional[list[int]] = None,
+    verbose: bool = False,
+) -> list[tree_utils.PyTree]:
+    if include_elements is None:
+        include_elements = list(range(len(list_of_data)))
+
+    assert min(include_elements) >= 0
+    assert max(include_elements) < len(list_of_data)
+
+    def _is_nan(ele: tree_utils.PyTree, i: int):
+        isnan = np.any(
+            [np.any(np.isnan(arr)) for arr in jax.tree_util.tree_leaves(ele)]
+        )
+        if isnan:
+            if verbose:
+                print(f"Sample with idx={i} is nan. It will be replaced.")
+            return True
+        return False
+
+    list_of_data_nonan = []
+    for i, ele in enumerate(list_of_data):
+        if _is_nan(ele, i):
+            while True:
+                j = random.choice(include_elements)
+                if not _is_nan(list_of_data[j], j):
+                    ele = list_of_data[j]
+                    break
+        list_of_data_nonan.append(ele)
+    return list_of_data_nonan