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

ring/algorithms/generator/batch.py

@@ -1,14 +1,12 @@
 from pathlib import Path
 import random
 from typing import Optional
-import warnings
 
 import jax
 import jax.numpy as jnp
 import numpy as np
 from tqdm import tqdm
 import tree_utils
-from tree_utils import tree_batch
 
 from ring import utils
 from ring.algorithms.generator import types
@@ -21,19 +19,13 @@ def _build_batch_matrix(batchsizes: list[int]) -> jax.Array:
     return jnp.array(arr)
 
 
-def batch_generators_lazy(
-    generators: types.Generator | list[types.Generator],
-    batchsizes: int | list[int] = 1,
+def generators_lazy(
+    generators: list[types.BatchedGenerator],
+    repeats: list[int],
     jit: bool = True,
 ) -> types.BatchedGenerator:
-    """Create a large generator by stacking multiple generators lazily."""
-    generators = utils.to_list(generators)
 
-    generators, batchsizes = _process_sizes_batchsizes_generators(
-        generators, batchsizes
-    )
-
-    batch_arr = _build_batch_matrix(batchsizes)
+    batch_arr = _build_batch_matrix(repeats)
     bs_total = len(batch_arr)
     pmap, vmap = utils.distribute_batchsize(bs_total)
     batch_arr = batch_arr.reshape((pmap, vmap))
@@ -56,60 +48,43 @@ def batch_generators_lazy(
         data = _generator(pmap_vmap_keys, batch_arr)
 
         # merge pmap and vmap axis
-        data = utils.merge_batchsize(data, pmap, vmap)
+        data = utils.merge_batchsize(data, pmap, vmap, third_dim_also=True)
         return data
 
     return generator
 
 
-def _number_of_executions_required(size: int) -> int:
-    _, vmap = utils.distribute_batchsize(size)
-
-    eager_threshold = utils.batchsize_thresholds()[1]
-    primes = iter(utils.primes(vmap))
-    n_calls = 1
-    while vmap > eager_threshold:
-        prime = next(primes)
-        n_calls *= prime
-        vmap /= prime
-
-    return n_calls
-
-
-def batch_generators_eager_to_list(
-    generators: types.Generator | list[types.Generator],
-    sizes: int | list[int],
+def generators_eager_to_list(
+    generators: list[types.BatchedGenerator],
+    n_calls: list[int],
     seed: int = 1,
+    disable_tqdm: bool = False,
 ) -> list[tree_utils.PyTree]:
-    "Returns list of unbatched sequences as numpy arrays."
-    generators, sizes = _process_sizes_batchsizes_generators(generators, sizes)
 
     key = jax.random.PRNGKey(seed)
     data = []
-    for gen, size in tqdm(
-        zip(generators, sizes),
+    for gen, n_call in tqdm(
+        zip(generators, n_calls),
         desc="executing generators",
-        total=len(sizes),
+        total=len(generators),
+        disable=disable_tqdm,
     ):
-
-        n_calls = _number_of_executions_required(size)
-        # decrease size by n_calls times
-        size = int(size / n_calls)
-        jit = True if n_calls > 1 else False
-        gen_jit = batch_generators_lazy(gen, size, jit=jit)
-
         for _ in tqdm(
-            range(n_calls),
+            range(n_call),
             desc="number of calls for each generator",
-            total=n_calls,
+            total=n_call,
             leave=False,
+            disable=disable_tqdm,
         ):
             key, consume = jax.random.split(key)
-            sample = gen_jit(consume)
+            sample = gen(consume)
             # converts also to numpy; but with np.array.flags.writeable = False
             sample = jax.device_get(sample)
             # this then sets this flag to True
             sample = jax.tree_map(np.array, sample)
+
+            sample_flat, _ = jax.tree_util.tree_flatten(sample)
+            size = 1 if len(sample_flat) == 0 else sample_flat[0].shape[0]
             data.extend([jax.tree_map(lambda a: a[i], sample) for i in range(size)])
 
     return data
@@ -197,7 +172,9 @@ def _data_fn_from_paths(
             ]
 
         def data_fn(indices: list[int]):
-            return tree_batch([list_of_data[i] for i in indices], backend="numpy")
+            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))
@@ -205,7 +182,7 @@ def _data_fn_from_paths(
     return data_fn, include_samples.copy()
 
 
-def _generator_from_data_fn(
+def generator_from_data_fn(
     data_fn,
     include_samples: list[int],
     shuffle: bool,
@@ -231,7 +208,7 @@ def _generator_from_data_fn(
     return generator
 
 
-def batched_generator_from_paths(
+def generator_from_paths(
     paths: list[str],
     batchsize: int,
     include_samples: Optional[list[int]] = None,
@@ -247,66 +224,6 @@ def batched_generator_from_paths(
     N = len(include_samples)
     assert N >= batchsize
 
-    generator = _generator_from_data_fn(data_fn, include_samples, shuffle, batchsize)
+    generator = generator_from_data_fn(data_fn, include_samples, shuffle, batchsize)
 
     return generator, N
-
-
-def batched_generator_from_list(
-    data: list,
-    batchsize: int,
-    shuffle: bool = True,
-    drop_last: bool = True,
-) -> types.BatchedGenerator:
-    assert drop_last, "Not `drop_last` is currently not implemented."
-    assert len(data) >= batchsize
-
-    def data_fn(indices: list[int]):
-        return tree_batch([data[i] for i in indices])
-
-    return _generator_from_data_fn(data_fn, list(range(len(data))), shuffle, batchsize)
-
-
-def batch_generators_eager(
-    generators: types.Generator | list[types.Generator],
-    sizes: int | list[int],
-    batchsize: int,
-    shuffle: bool = True,
-    drop_last: bool = True,
-    seed: int = 1,
-) -> types.BatchedGenerator:
-    """Eagerly create a large precomputed generator by calling multiple generators
-    and stacking their output."""
-
-    data = batch_generators_eager_to_list(generators, sizes, seed=seed)
-    return batched_generator_from_list(data, batchsize, shuffle, drop_last)
-
-
-def _process_sizes_batchsizes_generators(
-    generators: types.Generator | list[types.Generator],
-    batchsizes_or_sizes: int | list[int],
-) -> tuple[list, list]:
-    generators = utils.to_list(generators)
-    assert len(generators) > 0, "No generator was passed."
-
-    if isinstance(batchsizes_or_sizes, int):
-        assert (
-            batchsizes_or_sizes // len(generators)
-        ) > 0, f"Batchsize or size too small. {batchsizes_or_sizes} < {len(generators)}"
-        list_sizes = len(generators) * [batchsizes_or_sizes // len(generators)]
-    else:
-        list_sizes = batchsizes_or_sizes
-        assert 0 not in list_sizes
-
-    assert len(generators) == len(list_sizes)
-
-    _WARN_SIZE = 1e6  # disable this warning
-    for size in list_sizes:
-        if size >= _WARN_SIZE:
-            warnings.warn(
-                f"A generator will be called with a large batchsize of {size} "
-                f"(warn limit is {_WARN_SIZE}). The generator sizes are {list_sizes}."
-            )
-            break
-
-    return generators, list_sizes

ring/algorithms/generator/finalize_fns.py

@@ -0,0 +1,306 @@
+from typing import Optional
+import warnings
+
+import jax
+import jax.numpy as jnp
+import numpy as np
+import tree_utils
+
+from ring import base
+from ring import utils
+from ring.algorithms import sensors
+from ring.algorithms.generator import pd_control
+from ring.algorithms.generator import types
+
+
+class FinalizeFn:
+    def __init__(self, finalize_fn: types.FINALIZE_FN):
+        self.finalize_fn = finalize_fn
+
+    def __call__(self, Xy, extras):
+        (X, y), (key, *extras) = Xy, extras
+        # make sure we aren't overwriting anything
+        assert len(X) == len(y) == 0, f"X.keys={X.keys()}, y.keys={y.keys()}"
+        key, consume = jax.random.split(key)
+        Xy = self.finalize_fn(consume, *extras)
+        return Xy, tuple([key] + extras)
+
+
+def _rename_links(d: dict[str, dict], names: list[str]) -> dict[int, dict]:
+    for key in list(d.keys()):
+        if key in names:
+            d[str(names.index(key))] = d.pop(key)
+        else:
+            warnings.warn(
+                f"The key `{key}` was not found in names `{names}`. "
+                "It will not be renamed."
+            )
+
+    return d
+
+
+class Names2Indices:
+    def __init__(self, sys_noimu: base.System) -> None:
+        self.sys_noimu = sys_noimu
+
+    def __call__(self, Xy, extras):
+        (X, y), extras = Xy, extras
+        X = _rename_links(X, self.sys_noimu.link_names)
+        y = _rename_links(y, self.sys_noimu.link_names)
+        return (X, y), extras
+
+
+class JointAxisSensor:
+    def __init__(self, sys: base.System, **kwargs):
+        self.sys = sys
+        self.kwargs = kwargs
+
+    def __call__(self, Xy, extras):
+        (X, y), (key, q, x, sys_x) = Xy, extras
+        key, consume = jax.random.split(key)
+        X_joint_axes = sensors.joint_axes(
+            self.sys, x, sys_x, key=consume, **self.kwargs
+        )
+        X = utils.dict_union(X, X_joint_axes)
+        return (X, y), (key, q, x, sys_x)
+
+
+class RelPose:
+    def __init__(self, sys: base.System):
+        self.sys = sys
+
+    def __call__(self, Xy, extras):
+        (X, y), (key, q, x, sys_x) = Xy, extras
+        y_relpose = sensors.rel_pose(self.sys, x, sys_x)
+        y = utils.dict_union(y, y_relpose)
+        return (X, y), (key, q, x, sys_x)
+
+
+class RootIncl:
+    def __init__(self, sys: base.System):
+        self.sys = sys
+
+    def __call__(self, Xy, extras):
+        (X, y), (key, q, x, sys_x) = Xy, extras
+        y_root_incl = sensors.root_incl(self.sys, x, sys_x)
+        y = utils.dict_union(y, y_root_incl)
+        return (X, y), (key, q, x, sys_x)
+
+
+_default_imu_kwargs = dict(
+    noisy=True,
+    low_pass_filter_pos_f_cutoff=13.5,
+    low_pass_filter_rot_cutoff=16.0,
+)
+
+
+class IMU:
+    def __init__(self, **imu_kwargs):
+        self.kwargs = _default_imu_kwargs.copy()
+        self.kwargs.update(imu_kwargs)
+
+    def __call__(self, Xy: types.Xy, extras: types.OutputExtras):
+        (X, y), (key, q, x, sys) = Xy, extras
+        key, consume = jax.random.split(key)
+        X_imu = _imu_data(consume, x, sys, **self.kwargs)
+        X = utils.dict_union(X, X_imu)
+        return (X, y), (key, q, x, sys)
+
+
+def _imu_data(key, xs, sys_xs, **kwargs) -> dict:
+    sys_noimu, imu_attachment = sys_xs.make_sys_noimu()
+    inv_imu_attachment = {val: key for key, val in imu_attachment.items()}
+    X = {}
+    N = xs.shape()
+    for segment in sys_noimu.link_names:
+        if segment in inv_imu_attachment:
+            imu = inv_imu_attachment[segment]
+            key, consume = jax.random.split(key)
+            imu_measurements = sensors.imu(
+                xs=xs.take(sys_xs.name_to_idx(imu), 1),
+                gravity=sys_xs.gravity,
+                dt=sys_xs.dt,
+                key=consume,
+                **kwargs,
+            )
+        else:
+            imu_measurements = {
+                "acc": jnp.zeros(
+                    (
+                        N,
+                        3,
+                    )
+                ),
+                "gyr": jnp.zeros(
+                    (
+                        N,
+                        3,
+                    )
+                ),
+            }
+        X[segment] = imu_measurements
+    return X
+
+
+P_rot, P_pos = 100.0, 250.0
+_P_gains = {
+    "free": jnp.array(3 * [P_rot] + 3 * [P_pos]),
+    "free_2d": jnp.array(1 * [P_rot] + 2 * [P_pos]),
+    "px": jnp.array([P_pos]),
+    "py": jnp.array([P_pos]),
+    "pz": jnp.array([P_pos]),
+    "rx": jnp.array([P_rot]),
+    "ry": jnp.array([P_rot]),
+    "rz": jnp.array([P_rot]),
+    "rr": jnp.array([P_rot]),
+    # primary, residual
+    "rr_imp": jnp.array([P_rot, P_rot]),
+    "cor": jnp.array(3 * [P_rot] + 6 * [P_pos]),
+    "spherical": jnp.array(3 * [P_rot]),
+    "p3d": jnp.array(3 * [P_pos]),
+    "saddle": jnp.array([P_rot, P_rot]),
+    "frozen": jnp.array([]),
+    "suntay": jnp.array([P_rot]),
+}
+
+
+class DynamicalSimulation:
+    def __init__(
+        self,
+        custom_P_gains: dict[str, jax.Array] = dict(),
+        unactuated_subsystems: list[str] = [],
+        return_q_ref: bool = False,
+        overwrite_q_ref: Optional[tuple[jax.Array, dict[str, slice]]] = None,
+        **unroll_kwargs,
+    ):
+        self.unactuated_links = unactuated_subsystems
+        self.custom_P_gains = custom_P_gains
+        self.return_q_ref = return_q_ref
+        self.overwrite_q_ref = overwrite_q_ref
+        self.unroll_kwargs = unroll_kwargs
+
+    def __call__(
+        self, Xy: types.Xy, extras: types.OutputExtras
+    ) -> tuple[types.Xy, types.OutputExtras]:
+        (X, y), (key, q, _, sys_x) = Xy, extras
+        idx_map_q = sys_x.idx_map("q")
+
+        if self.overwrite_q_ref is not None:
+            q, idx_map_q = self.overwrite_q_ref
+            assert q.shape[-1] == sum([s.stop - s.start for s in idx_map_q.values()])
+
+        sys_q_ref = sys_x
+        if len(self.unactuated_links) > 0:
+            sys_q_ref = sys_x.delete_system(self.unactuated_links)
+
+        q_ref = []
+        p_gains_list = []
+        q = q.T
+
+        def build_q_ref(_, __, name, link_type):
+            q_ref.append(q[idx_map_q[name]])
+
+            if link_type in self.custom_P_gains:
+                p_gain_this_link = self.custom_P_gains[link_type]
+            elif link_type in _P_gains:
+                p_gain_this_link = _P_gains[link_type]
+            else:
+                raise RuntimeError(
+                    f"Please proved gain parameters for the joint typ `{link_type}`"
+                    " via the argument `custom_P_gains: dict[str, Array]`"
+                )
+
+            required_qd_size = base.QD_WIDTHS[link_type]
+            assert (
+                required_qd_size == p_gain_this_link.size
+            ), f"The gain parameters must be of qd_size=`{required_qd_size}`"
+            f" but got `{p_gain_this_link.size}`. This happened for the link "
+            f"`{name}` of type `{link_type}`."
+            p_gains_list.append(p_gain_this_link)
+
+        sys_q_ref.scan(build_q_ref, "ll", sys_q_ref.link_names, sys_q_ref.link_types)
+        q_ref, p_gains_array = jnp.concatenate(q_ref).T, jnp.concatenate(p_gains_list)
+
+        # perform dynamical simulation
+        states = pd_control._unroll_dynamics_pd_control(
+            sys_x, q_ref, p_gains_array, sys_q_ref=sys_q_ref, **self.unroll_kwargs
+        )
+
+        if self.return_q_ref:
+            X = utils.dict_union(X, dict(q_ref=q_ref))
+
+        return (X, y), (key, states.q, states.x, sys_x)
+
+
+def _flatten(seq: list):
+    seq = tree_utils.tree_batch(seq, backend=None)
+    seq = tree_utils.batch_concat_acme(seq, num_batch_dims=3).transpose((1, 2, 0, 3))
+    return seq
+
+
+def _expand_dt(X: dict, T: int):
+    dt = X.pop("dt", None)
+    if dt is not None:
+        if isinstance(dt, np.ndarray):
+            numpy = np
+        else:
+            numpy = jnp
+        dt = numpy.repeat(dt[:, None, :], T, axis=1)
+        for seg in X:
+            X[seg]["dt"] = dt
+    return X
+
+
+def _expand_then_flatten(args):
+    X, y = args
+    gyr = X["0"]["gyr"]
+
+    batched = True
+    if gyr.ndim == 2:
+        batched = False
+        X, y = tree_utils.add_batch_dim((X, y))
+
+    X = _expand_dt(X, gyr.shape[-2])
+
+    N = len(X)
+
+    def dict_to_tuple(d: dict[str, jax.Array]):
+        tup = (d["acc"], d["gyr"])
+        if "joint_axes" in d:
+            tup = tup + (d["joint_axes"],)
+        if "dt" in d:
+            tup = tup + (d["dt"],)
+        return tup
+
+    X = [dict_to_tuple(X[str(i)]) for i in range(N)]
+    y = [y[str(i)] for i in range(N)]
+
+    X, y = _flatten(X), _flatten(y)
+    if not batched:
+        X, y = jax.tree_map(lambda arr: arr[0], (X, y))
+    return X, y
+
+
+class GeneratorTrafoLambda:
+    def __init__(self, f, input: bool = False):
+        self.f = f
+        self.input = input
+
+    def __call__(self, gen):
+        if self.input:
+
+            def _gen(*args):
+                return gen(*self.f(*args))
+
+        else:
+
+            def _gen(*args):
+                return self.f(gen(*args))
+
+        return _gen
+
+
+def GeneratorTrafoExpandFlatten(gen, jit: bool = False):
+    if jit:
+        return GeneratorTrafoLambda(jax.jit(_expand_then_flatten))(gen)
+    return GeneratorTrafoLambda(_expand_then_flatten)(gen)

ring/algorithms/generator/motion_artifacts.py

@@ -1,3 +1,4 @@
+import inspect
 import warnings
 
 import jax
@@ -50,7 +51,8 @@ def inject_subsystems(
     translational_stif: float = 50.0,
     translational_damp: float = 0.1,
     disable_warning: bool = False,
-    **kwargs,
+    **kwargs,  # needed because `imu_motion_artifacts_kwargs` is used
+    # for `setup_fn_randomize_damping_stiffness_factory` also
 ) -> base.System:
     imu_idx_to_name_map = {sys.name_to_idx(imu): imu for imu in sys.findall_imus()}
 
@@ -123,9 +125,10 @@ def _log_uniform(key, shape, minval, maxval):
 
 
 def setup_fn_randomize_damping_stiffness_factory(
-    prob_rigid: float,
-    all_imus_either_rigid_or_flex: bool,
-    imus_surely_rigid: list[str],
+    prob_rigid: float = 0.0,
+    all_imus_either_rigid_or_flex: bool = False,
+    imus_surely_rigid: list[str] = [],
+    **kwargs,
 ):
     assert 0 <= prob_rigid <= 1
     assert prob_rigid != 1, "Use `imu_motion_artifacts`=False instead."
@@ -197,6 +200,18 @@ def setup_fn_randomize_damping_stiffness_factory(
     return setup_fn_randomize_damping_stiffness
 
 
+# assert that there exists no keyword arg duplicate which would induce ambiguity
+kwargs = lambda f: set(inspect.signature(f).parameters.keys())
+assert (
+    len(
+        kwargs(inject_subsystems).intersection(
+            kwargs(setup_fn_randomize_damping_stiffness_factory)
+        )
+    )
+    == 1
+)
+
+
 def _match_q_x_between_sys(
     sys_small: base.System,
     q_large: jax.Array,
@@ -228,21 +243,18 @@ def _match_q_x_between_sys(
     return q_small, x_small
 
 
-class GeneratorTrafoHideInjectedBodies:
-    def __call__(self, gen):
-        def _gen(*args):
-            (X, y), (key, q, x, sys_x) = gen(*args)
-
-            # delete injected frames; then rename from `_imu` back to `imu`
-            imus = sys_x.findall_imus()
-            _imu2imu_map = {imu_reference_link_name(imu): imu for imu in imus}
-            sys = sys_x.delete_system(imus)
-            for _imu, imu in _imu2imu_map.items():
-                sys = sys.change_link_name(_imu, imu).change_joint_type(imu, "frozen")
+class HideInjectedBodies:
+    def __call__(self, Xy, extras):
+        (X, y), (key, q, x, sys_x) = Xy, extras
 
-            # match q and x to `sys`; second axis is link axis
-            q, x = _match_q_x_between_sys(sys, q, x, sys_x, q_large_skip=imus)
+        # delete injected frames; then rename from `_imu` back to `imu`
+        imus = sys_x.findall_imus()
+        _imu2imu_map = {imu_reference_link_name(imu): imu for imu in imus}
+        sys = sys_x.delete_system(imus)
+        for _imu, imu in _imu2imu_map.items():
+            sys = sys.change_link_name(_imu, imu).change_joint_type(imu, "frozen")
 
-            return (X, y), (key, q, x, sys)
+        # match q and x to `sys`; second axis is link axis
+        q, x = _match_q_x_between_sys(sys, q, x, sys_x, q_large_skip=imus)
 
-        return _gen
+        return (X, y), (key, q, x, sys)

ring/algorithms/generator/setup_fns.py

@@ -0,0 +1,43 @@
+import jax
+import jax.numpy as jnp
+
+from ring import base
+from ring import maths
+
+
+def _setup_fn_randomize_positions(key: jax.Array, sys: base.System) -> base.System:
+    ts = sys.links.transform1
+
+    for i in range(sys.num_links()):
+        link = sys.links[i]
+        key, new_pos = _draw_pos_uniform(key, link.pos_min, link.pos_max)
+        ts = ts.index_set(i, ts[i].replace(pos=new_pos))
+
+    return sys.replace(links=sys.links.replace(transform1=ts))
+
+
+def _draw_pos_uniform(key, pos_min, pos_max):
+    key, c1, c2, c3 = jax.random.split(key, num=4)
+    pos = jnp.array(
+        [
+            jax.random.uniform(c1, minval=pos_min[0], maxval=pos_max[0]),
+            jax.random.uniform(c2, minval=pos_min[1], maxval=pos_max[1]),
+            jax.random.uniform(c3, minval=pos_min[2], maxval=pos_max[2]),
+        ]
+    )
+    return key, pos
+
+
+def _setup_fn_randomize_transform1_rot(
+    key, sys, maxval: float, not_imus: bool = True
+) -> base.System:
+    new_transform1 = sys.links.transform1.replace(
+        rot=maths.quat_random(key, (sys.num_links(),), maxval=maxval)
+    )
+    if not_imus:
+        imus = [name for name in sys.link_names if name[:3] == "imu"]
+        new_rot = new_transform1.rot
+        for imu in imus:
+            new_rot = new_rot.at[sys.name_to_idx(imu)].set(jnp.array([1.0, 0, 0, 0]))
+        new_transform1 = new_transform1.replace(rot=new_rot)
+    return sys.replace(links=sys.links.replace(transform1=new_transform1))

ring/algorithms/generator/types.py

@@ -1,9 +1,10 @@
-from typing import Callable, Protocol
+from typing import Callable
 
 import jax
-from ring import base
 from tree_utils import PyTree
 
+from ring import base
+
 PRNGKey = jax.Array
 InputExtras = base.System
 OutputExtras = tuple[PRNGKey, jax.Array, jax.Array, base.System]
@@ -18,19 +19,3 @@ Generator = Callable[[PRNGKey], Xy]
 BatchedGenerator = Callable[[PRNGKey], BatchedXy]
 SETUP_FN = Callable[[PRNGKey, base.System], base.System]
 FINALIZE_FN = Callable[[PRNGKey, jax.Array, base.Transform, base.System], Xy]
-
-
-class GeneratorTrafo(Protocol):
-    def __call__(  # noqa: E704
-        self,
-        gen: (
-            GeneratorWithInputOutputExtras
-            | GeneratorWithOutputExtras
-            | GeneratorWithInputExtras
-        ),
-    ) -> (
-        GeneratorWithInputOutputExtras
-        | GeneratorWithOutputExtras
-        | GeneratorWithInputExtras
-        | Generator
-    ): ...