imt-ring 1.6.46__py3-none-any.whl → 1.7.0__py3-none-any.whl

{imt_ring-1.6.46.dist-info → imt_ring-1.7.0.dist-info}/METADATA

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

@@ -1,6 +1,6 @@
-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
@@ -15,13 +15,25 @@ ring/algorithms/custom_joints/rr_joint.py,sha256=jnRtjtOCALMaq2_0bcu2d7qgfQ6etXp
 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=vQEzcBDdG3BPqTGEktC74DsCfvgKktj9DKWK8gBzRtE,3805
+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.46.dist-info/METADATA,sha256=9NOkzI2PpdcJpw90_ZV0smHEWyhkrZyEF211Wy1gNpg,5888
-imt_ring-1.6.46.dist-info/WHEEL,sha256=jB7zZ3N9hIM9adW7qlTAyycLYW9npaWKLRzaoVcLKcM,91
-imt_ring-1.6.46.dist-info/entry_points.txt,sha256=npNqSOvNiBR0BNa_GL3J66q8Gky3h0G_PHzHzk8oyE0,66
-imt_ring-1.6.46.dist-info/top_level.txt,sha256=EiT790-lAyi8iwTzJArH3f2k77rwhDn00q-4PlmvDQo,5
-imt_ring-1.6.46.dist-info/RECORD,,
+imt_ring-1.7.0.dist-info/METADATA,sha256=CNwgvWr9Yu7MgIfcNwXkuByr7_8vkxvO5IkJg3iDKbs,5887
+imt_ring-1.7.0.dist-info/WHEEL,sha256=SmOxYU7pzNKBqASvQJ7DjX3XGUF92lrGhMb3R6_iiqI,91
+imt_ring-1.7.0.dist-info/entry_points.txt,sha256=npNqSOvNiBR0BNa_GL3J66q8Gky3h0G_PHzHzk8oyE0,66
+imt_ring-1.7.0.dist-info/top_level.txt,sha256=EiT790-lAyi8iwTzJArH3f2k77rwhDn00q-4PlmvDQo,5
+imt_ring-1.7.0.dist-info/RECORD,,

{imt_ring-1.6.46.dist-info → imt_ring-1.7.0.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/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/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()
+    dt_target = sys.dt
+
+    with InteractiveViewer(sys, **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