imt-ring 1.2.1__py3-none-any.whl

ring/rendering/base_render.py

@@ -0,0 +1,271 @@
+from typing import Optional
+
+import jax
+import jax.numpy as jnp
+import numpy as np
+import tqdm
+
+from ring import algebra
+from ring import base
+from ring import maths
+from ring import sim2real
+from ring import utils
+from ring.algorithms import kinematics
+
+_rgbas = {
+    "self": (0.7, 0.5, 0.1, 1.0),
+    "effector": (0.7, 0.4, 0.1, 1.0),
+    "decoration": (0.3, 0.5, 0.7, 1.0),
+    "eye": (0.0, 0.2, 1.0, 1.0),
+    "target": (0.6, 0.3, 0.3, 1.0),
+    "site": (0.5, 0.5, 0.5, 0.3),
+    "red": (0.8, 0.2, 0.2, 1.0),
+    "green": (0.2, 0.8, 0.2, 1.0),
+    "blue": (0.2, 0.2, 0.8, 1.0),
+    "yellow": (0.8, 0.8, 0.2, 1.0),
+    "cyan": (0.2, 0.8, 0.8, 1.0),
+    "magenta": (0.8, 0.2, 0.8, 1.0),
+    "white": (0.8, 0.8, 0.8, 1.0),
+    "gray": (0.5, 0.5, 0.5, 1.0),
+    "brown": (0.6, 0.3, 0.1, 1.0),
+    "orange": (0.8, 0.5, 0.2, 1.0),
+    "pink": (0.8, 0.75, 0.8, 1.0),
+    "purple": (0.5, 0.2, 0.5, 1.0),
+    "lime": (0.5, 0.8, 0.2, 1.0),
+    "gold": (0.8, 0.84, 0.2, 1.0),
+    "matplotlib_green": (0.0, 0.502, 0.0, 1.0),
+    "matplotlib_blue": (0.012, 0.263, 0.8745, 1.0),
+    "matplotlib_lightblue": (0.482, 0.784, 0.9647, 1.0),
+    "matplotlib_salmon": (0.98, 0.502, 0.447, 1.0),
+    "black": (0.1, 0.1, 0.1, 1.0),
+    "dustin_exp_blue": (75 / 255, 93 / 255, 208 / 255, 1.0),
+    "dustin_exp_white": (241 / 255, 239 / 255, 208 / 255, 1.0),
+    "dustin_exp_orange": (227 / 255, 139 / 255, 61 / 255, 1.0),
+}
+
+
+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 backend == "mujoco":
+        utils.import_lib("mujoco")
+        from ring.rendering.mujoco_render import MujocoScene
+
+        scene = MujocoScene(**scene_kwargs)
+    elif backend == "vispy":
+        vispy = utils.import_lib("vispy")
+
+        if "vispy_backend" in scene_kwargs:
+            vispy_backend = scene_kwargs.pop("vispy_backend")
+        else:
+            vispy_backend = "pyqt6"
+
+        vispy.use(vispy_backend)
+
+        from ring.rendering.vispy_render import VispyScene
+
+        scene = VispyScene(**scene_kwargs)
+    else:
+        raise NotImplementedError
+
+    # mujoco does not implement the xyz Geometry; instead replace it with
+    # three capsule Geometries
+    geoms = sys.geoms
+    if backend == "mujoco":
+        geoms = _replace_xyz_geoms(geoms)
+
+    # convert all colors to rgbas
+    geoms_rgba = [_color_to_rgba(geom) for geom in geoms]
+
+    if xs is None:
+        xs = kinematics.forward_kinematics(sys, base.State.create(sys))[1].x
+
+    # convert time-axis of batched xs object into a list of unbatched x objects
+    if isinstance(xs, base.Transform) and xs.ndim() == 3:
+        xs = [xs[t] for t in range(xs.shape())]
+
+    # ensure that a single unbatched x object is also a list
+    xs = utils.to_list(xs)
+
+    if render_every_nth != 1:
+        xs = [xs[t] for t in range(0, len(xs), render_every_nth)]
+
+    n_links = sys.num_links()
+
+    def data_check(x):
+        assert (
+            x.pos.ndim == x.rot.ndim == 2
+        ), f"Expected shape = (n_links, 3/4). Got pos.shape{x.pos.shape}, "
+        "rot.shape={x.rot.shape}"
+        assert (
+            x.pos.shape[0] == x.rot.shape[0] == n_links
+        ), "Number of links does not match"
+
+    for x in xs:
+        data_check(x)
+
+    scene.init(geoms_rgba)
+
+    frames = []
+    for x in tqdm.tqdm(xs, "Rendering frames..", disable=not show_pbar):
+        scene.update(x)
+        frames.append(scene.render(camera=camera))
+
+    return frames
+
+
+def render_prediction(
+    sys: base.System,
+    xs: base.Transform | list[base.Transform],
+    yhat: dict | jax.Array | np.ndarray,
+    stepframe: int = 1,
+    # 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."
+    if isinstance(xs, list):
+        # list -> batched Transform
+        xs = xs[0].batch(*xs[1:])
+
+    sys_noimu, _ = sys.make_sys_noimu()
+
+    if isinstance(yhat, (np.ndarray, jax.Array)):
+        yhat = {name: yhat[..., i, :] for i, name in enumerate(sys_noimu.link_names)}
+
+    xs_noimu = sim2real.match_xs(sys_noimu, xs, sys)
+
+    # `yhat` are child-to-parent transforms, but we need parent-to-child
+    # but not for those that connect to root, those are already parent-to-child
+    transform2hat_rot = {}
+    for name, p in zip(sys_noimu.link_names, sys_noimu.link_parents):
+        if p == -1:
+            transform2hat_rot[name] = yhat[name]
+        else:
+            transform2hat_rot[name] = maths.quat_inv(yhat[name])
+
+    transform1, transform2 = sim2real.unzip_xs(sys_noimu, xs_noimu)
+
+    # we add the missing links in transform2hat, links that connect to worldbody
+    transform2hat = []
+    for i, name in enumerate(sys_noimu.link_names):
+        if name in transform2hat_rot:
+            transform2_name = base.Transform.create(rot=transform2hat_rot[name])
+        else:
+            transform2_name = transform2.take(i, axis=1)
+        transform2hat.append(transform2_name)
+
+    # after transpose shape is (n_timesteps, n_links, ...)
+    transform2hat = transform2hat[0].batch(*transform2hat[1:]).transpose((1, 0, 2))
+
+    xshat = sim2real.zip_xs(sys_noimu, transform1, transform2hat)
+
+    # swap time axis, and link axis
+    xs, xshat = xs.transpose((1, 0, 2)), xshat.transpose((1, 0, 2))
+    # create mapping from `name` -> Transform
+    xs_dict = dict(
+        zip(
+            ["hat_" + name for name in sys_noimu.link_names],
+            [xshat[i] for i in range(sys_noimu.num_links())],
+        )
+    )
+    xs_dict.update(
+        dict(
+            zip(
+                sys.link_names,
+                [xs[i] for i in range(sys.num_links())],
+            )
+        )
+    )
+
+    sys_render = _sys_render(sys, transparent_segment_to_root)
+    xs_render = []
+    for name in sys_render.link_names:
+        xs_render.append(xs_dict[name])
+    xs_render = xs_render[0].batch(*xs_render[1:])
+    xs_render = xs_render.transpose((1, 0, 2))
+    N = xs_render.shape()
+    xs_render = [xs_render[t] for t in range(0, N, stepframe)]
+
+    frames = render(sys_render, xs_render, **kwargs)
+
+    return frames
+
+
+def _color_to_rgba(geom: base.Geometry) -> base.Geometry:
+    if geom.color is None:
+        new_color = _rgbas["self"]
+    elif isinstance(geom.color, tuple):
+        if len(geom.color) == 3:
+            new_color = geom.color + (1.0,)
+        else:
+            new_color = geom.color
+    elif isinstance(geom.color, str):
+        new_color = _rgbas[geom.color]
+    else:
+        raise NotImplementedError
+
+    return geom.replace(color=new_color)
+
+
+def _xyz_to_three_capsules(xyz: base.XYZ) -> list[base.Geometry]:
+    capsules = []
+    length = xyz.size
+    radius = length / 6
+    colors = ["red", "green", "blue"]
+    rot_axis = [1, 0, 2]
+
+    for i, (color, axis) in enumerate(zip(colors, rot_axis)):
+        pos = maths.unit_vectors(i) * length / 2
+        rot = maths.quat_rot_axis(maths.unit_vectors(axis), jnp.pi / 2)
+        t = algebra.transform_mul(base.Transform(pos, rot), xyz.transform)
+        capsules.append(
+            base.Capsule(0.0, t, xyz.link_idx, color, xyz.edge_color, radius, length)
+        )
+    return capsules
+
+
+def _replace_xyz_geoms(geoms: list[base.Geometry]) -> list[base.Geometry]:
+    geoms_replaced = []
+    for geom in geoms:
+        if isinstance(geom, base.XYZ):
+            geoms_replaced += _xyz_to_three_capsules(geom)
+        else:
+            geoms_replaced.append(geom)
+    return geoms_replaced
+
+
+def _sys_render(
+    sys: base.Transform, transparent_segment_to_root: bool
+) -> base.Transform:
+    sys_noimu, _ = sys.make_sys_noimu()
+
+    def _geoms_replace_color(sys: base.System, color):
+        keep = lambda i: (not transparent_segment_to_root) or sys.link_parents[i] != -1
+        geoms = [g.replace(color=color) for g in sys.geoms if keep(g.link_idx)]
+        return sys.replace(geoms=geoms)
+
+    # replace render color of geoms for render of predicted motion
+    prediction_color = (78 / 255, 163 / 255, 243 / 255, 1.0)
+    sys_newcolor = _geoms_replace_color(sys_noimu, prediction_color)
+    sys_render = sys.inject_system(sys_newcolor.add_prefix_suffix("hat_"))
+
+    return sys_render

ring/rendering/mujoco_render.py

@@ -0,0 +1,222 @@
+from typing import Optional, Sequence
+
+import mujoco
+import numpy as np
+from ring import base
+from ring import maths
+
+_skybox = """<texture name="skybox" type="skybox" builtin="gradient" rgb1=".4 .6 .8" rgb2="0 0 0" width="800" height="800" mark="random" markrgb="1 1 1"/>"""  # noqa: E501
+_floor = """<geom name="floor" pos="0 0 -0.5" size="0 0 1" type="plane" material="matplane" mass="0"/>"""  # noqa: E501
+
+
+def _build_model_of_geoms(
+    geoms: list[base.Geometry],
+    cameras: dict[int, Sequence[str]],
+    lights: dict[int, Sequence[str]],
+    floor: bool,
+    stars: bool,
+    debug: bool,
+) -> mujoco.MjModel:
+    # sort in ascending order, this shouldn't be required as it is already done by
+    geoms = geoms.copy()
+    geoms.sort(key=lambda ele: ele.link_idx)
+
+    # range of required link_indices to which geoms attach
+    unique_parents = set([geom.link_idx for geom in geoms])
+
+    # throw error if you attached a camera or light to a body that has no geoms
+    inside_worldbody_cameras = ""
+    for camera_parent in cameras:
+        if -1 not in unique_parents:
+            if camera_parent == -1:
+                for camera_str in cameras[camera_parent]:
+                    inside_worldbody_cameras += camera_str
+                continue
+
+        assert (
+            camera_parent in unique_parents
+        ), f"Camera parent {camera_parent} not in {unique_parents}"
+
+    inside_worldbody_lights = ""
+    for light_parent in lights:
+        if -1 not in unique_parents:
+            if light_parent == -1:
+                for light_str in lights[light_parent]:
+                    inside_worldbody_lights += light_str
+                continue
+
+        assert (
+            light_parent in unique_parents
+        ), f"Light parent {light_parent} not in {unique_parents}"
+
+    # group together all geoms in each link
+    grouped_geoms = dict(
+        zip(unique_parents, [list() for _ in range(len(unique_parents))])
+    )
+    parent = -1
+    for geom in geoms:
+        while geom.link_idx != parent:
+            parent += 1
+        grouped_geoms[parent].append(geom)
+
+    inside_worldbody = ""
+    for parent, geoms in grouped_geoms.items():
+        find = lambda dic: dic[parent] if parent in dic else []
+        inside_worldbody += _xml_str_one_body(
+            parent, geoms, find(cameras), find(lights)
+        )
+
+    parents_noworld = unique_parents - set([-1])
+    targetbody = min(parents_noworld) if len(parents_noworld) > 0 else -1
+    xml_str = f""" # noqa: E501
+<mujoco>
+  <asset>
+    <texture name="texplane" type="2d" builtin="checker" rgb1=".25 .25 .25" rgb2=".3 .3 .3" width="512" height="512" mark="cross" markrgb=".8 .8 .8"/>
+    <material name="matplane" reflectance="0.3" texture="texplane" texrepeat="1 1" texuniform="true"/>
+    {_skybox if stars else ''}
+    <texture name="grid" type="2d" builtin="checker" rgb1=".1 .2 .3" rgb2=".2 .3 .4" width="300" height="300" mark="edge" markrgb=".2 .3 .4"/>
+    <material name="grid" texture="grid" texrepeat="1 1" texuniform="true" reflectance=".2"/>
+  </asset>
+
+  <visual>
+    <headlight ambient=".4 .4 .4" diffuse=".8 .8 .8" specular="0.1 0.1 0.1"/>
+    <map znear=".01"/>
+    <quality shadowsize="8192"/>
+    <global offwidth="3840" offheight="2160"/>
+  </visual>
+
+<worldbody>
+<camera pos="0 -1 1" name="trackcom" mode="trackcom"/>
+<camera pos="0 -1 1" name="target" mode="targetbodycom" target="{targetbody}"/>
+<camera pos="0 -3 3" name="targetfar" mode="targetbodycom" target="{targetbody}"/>
+<camera pos="0 -5 5" name="targetFar" mode="targetbodycom" target="{targetbody}"/>
+<light pos="0 0 4" dir="0 0 -1"/>
+{_floor if floor else ''}
+{inside_worldbody_cameras}
+{inside_worldbody_lights}
+{inside_worldbody}
+</worldbody>
+</mujoco>
+"""
+    if debug:
+        print("Mujoco xml string: ", xml_str)
+
+    return mujoco.MjModel.from_xml_string(xml_str)
+
+
+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)
+
+    inside_body_cameras = ""
+    for camera in cameras:
+        inside_body_cameras += camera  # + "\n"
+
+    inside_body_lights = ""
+    for light in lights:
+        inside_body_lights += light  # + "\n"
+
+    return f"""
+<body name="{body_number}" mocap="true">
+{inside_body_cameras}
+{inside_body_lights}
+{inside_body_geoms}
+</body>
+"""
+
+
+def _xml_str_one_geom(geom: base.Geometry) -> str:
+    rgba = f'rgba="{_array_to_str(geom.color)}"'
+
+    if isinstance(geom, base.Box):
+        type_size = f'type="box" size="{_array_to_str([geom.dim_x / 2, geom.dim_y / 2, geom.dim_z / 2])}"'  # noqa: E501
+    elif isinstance(geom, base.Sphere):
+        type_size = f'type="sphere" size="{_array_to_str([geom.radius])}"'
+    elif isinstance(geom, base.Capsule):
+        type_size = (
+            f'type="capsule" size="{_array_to_str([geom.radius, geom.length / 2])}"'
+        )
+    elif isinstance(geom, base.Cylinder):
+        type_size = (
+            f'type="cylinder" size="{_array_to_str([geom.radius, geom.length / 2])}"'
+        )
+    else:
+        raise NotImplementedError
+
+    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}/>"
+
+
+def _array_to_str(arr: Sequence[float]) -> str:
+    # TODO; remove round & truncation
+    return "".join(["{:.4f} ".format(np.round(value, 4)) for value in arr])[:-1]
+
+
+class MujocoScene:
+    def __init__(
+        self,
+        height: int = 240,
+        width: int = 320,
+        add_cameras: dict[int, str | Sequence[str]] = {},
+        add_lights: dict[int, str | Sequence[str]] = {},
+        show_stars: bool = True,
+        show_floor: bool = True,
+        debug: bool = False,
+    ) -> None:
+        self.debug = debug
+        self.height, self.width = height, width
+
+        def to_list(dic: dict):
+            for k, v in dic.items():
+                if isinstance(v, str):
+                    dic[k] = [v]
+            return dic
+
+        self.add_cameras, self.add_lights = to_list(add_cameras), to_list(add_lights)
+        self.show_stars = show_stars
+        self.show_floor = show_floor
+
+    def init(self, geoms: list[base.Geometry]):
+        self._parent_ids = list(set([geom.link_idx for geom in geoms]))
+        self._model = _build_model_of_geoms(
+            geoms,
+            self.add_cameras,
+            self.add_lights,
+            floor=self.show_floor,
+            stars=self.show_stars,
+            debug=self.debug,
+        )
+        self._data = mujoco.MjData(self._model)
+        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
+        for parent_id in self._parent_ids:
+            if parent_id == -1:
+                continue
+
+            # body name is just the str(parent_id)
+            # squeeze reduces shape (1,) to () which removes a warning
+            mocap_id = int(np.squeeze(self._model.body(str(parent_id)).mocapid))
+
+            if self.debug:
+                print(f"link_idx: {parent_id}, mocap_id: {mocap_id}")
+
+            mocap_pos = pos[parent_id]
+            mocap_quat = rot[parent_id]
+            self._data.mocap_pos[mocap_id] = mocap_pos
+            self._data.mocap_quat[mocap_id] = mocap_quat
+
+        if self.debug:
+            print("mocap_pos: ", self._data.mocap_pos)
+            print("mocap_quat: ", self._data.mocap_quat)
+
+        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)
+        return self._renderer.render()