imt-ring 1.5.2__py3-none-any.whl → 1.6.1__py3-none-any.whl

{imt_ring-1.5.2.dist-info → imt_ring-1.6.1.dist-info}/METADATA

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

@@ -1,6 +1,6 @@
 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
@@ -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
@@ -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=Y8B2V647JMM57S3GmCwAjCM4XuN5RwMLhcDfjReP3kQ,6526
-imt_ring-1.5.2.dist-info/METADATA,sha256=YhkKO-ToWNUrygQCGNFqn6Ugph4_ZVHdLK8W7LnL2n0,3104
-imt_ring-1.5.2.dist-info/WHEEL,sha256=Z4pYXqR_rTB7OWNDYFOm1qRk0RX6GFP2o8LgvP453Hk,91
-imt_ring-1.5.2.dist-info/top_level.txt,sha256=EiT790-lAyi8iwTzJArH3f2k77rwhDn00q-4PlmvDQo,5
-imt_ring-1.5.2.dist-info/RECORD,,
+ring/utils/utils.py,sha256=k7t-QxMWrNRnjfNB9rSobmLCmhJigE8__gkT-Il0Ee4,6492
+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.1.dist-info/METADATA,sha256=FrI0S7Njj9yZgqG5Wuek8KFocnUOV18c7Ar2T_V0ums,3104
+imt_ring-1.6.1.dist-info/WHEEL,sha256=Z4pYXqR_rTB7OWNDYFOm1qRk0RX6GFP2o8LgvP453Hk,91
+imt_ring-1.6.1.dist-info/top_level.txt,sha256=EiT790-lAyi8iwTzJArH3f2k77rwhDn00q-4PlmvDQo,5
+imt_ring-1.6.1.dist-info/RECORD,,

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/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/register_gym_envs/__init__.py

@@ -0,0 +1,3 @@
+import gymnasium as gym
+
+gym.register("Saddle-v0", "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

@@ -16,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):