imt-ring 1.6.37__py3-none-any.whl → 1.6.39__py3-none-any.whl

ring/rendering/mujoco_render.py

@@ -10,8 +10,8 @@ _skybox = """<texture name="skybox" type="skybox" builtin="gradient" rgb1=".4 .6
 _skybox_white = """<texture name="skybox" type="skybox" builtin="gradient" rgb1="1 1 1" rgb2="1 1 1" width="800" height="800" mark="random" markrgb="1 1 1"/>"""  # noqa: E501
 
 
-def _floor(floor_z: float) -> str:
-    return f"""<geom name="floor" pos="0 0 {floor_z}" size="0 0 1" type="plane" material="matplane" mass="0"/>"""  # noqa: E501
+def _floor(z: float, material: str) -> str:
+    return f"""<geom name="floor" pos="0 0 {z}" size="0 0 1" type="plane" material="{material}" mass="0"/>"""  # noqa: E501
 
 
 def _build_model_of_geoms(
@@ -19,7 +19,7 @@ def _build_model_of_geoms(
     cameras: dict[int, Sequence[str]],
     lights: dict[int, Sequence[str]],
     floor: bool,
-    floor_z: float,
+    floor_kwargs: dict,
     stars: bool,
     debug: bool,
 ) -> mujoco.MjModel:
@@ -77,10 +77,13 @@ def _build_model_of_geoms(
     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"/>
+    <texture name="texplane" type="2d" builtin="checker" rgb1=".25 .25 .25" rgb2=".3 .3 .3" width="512" height="512" mark="cross" markrgb=".3 .3 .3"/>
     <material name="matplane" reflectance="0.3" texture="texplane" texrepeat="1 1" texuniform="true"/>
     <texture type="2d" name="groundplane" builtin="checker" mark="edge" rgb1="0.2 0.3 0.4" rgb2="0.1 0.2 0.3" markrgb="0.8 0.8 0.8" width="300" height="300"/>
     <material name="groundplane" texture="groundplane" texuniform="true" texrepeat="2 2" reflectance="0.2"/>
+    <material name="beige" rgba="0.76 0.80 0.50 1.0" specular="0.3" shininess="0.1" />
+    <material name="white" rgba="0.9 0.9 0.9 1.0" reflectance="0"/>
+    <material name="gray" rgba="0.4 0.5 0.5 1.0" reflectance="0.25"/>
     {_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"/>
@@ -98,7 +101,7 @@ def _build_model_of_geoms(
 <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}"/>
-{_floor(floor_z) if floor else ''}
+{_floor(**floor_kwargs) if floor else ''}
 {inside_worldbody_cameras}
 {inside_worldbody_lights}
 {inside_worldbody}
@@ -176,6 +179,7 @@ class MujocoScene:
         show_stars: bool = True,
         show_floor: bool = True,
         floor_z: float = -0.84,
+        floor_material: str = "matplane",
         debug: bool = False,
     ) -> None:
         self.debug = debug
@@ -190,7 +194,7 @@ class MujocoScene:
         self.add_cameras, self.add_lights = to_list(add_cameras), to_list(add_lights)
         self.show_stars = show_stars
         self.show_floor = show_floor
-        self.floor_z = floor_z
+        self.floor_kwargs = dict(z=floor_z, material=floor_material)
 
     def init(self, geoms: list[base.Geometry]):
         self._parent_ids = list(set([geom.link_idx for geom in geoms]))
@@ -199,7 +203,7 @@ class MujocoScene:
             self.add_cameras,
             self.add_lights,
             floor=self.show_floor,
-            floor_z=self.floor_z,
+            floor_kwargs=self.floor_kwargs,
             stars=self.show_stars,
             debug=self.debug,
         )

ring/rendering/vispy_render.py

@@ -7,14 +7,15 @@ from typing import Optional, TypeVar
 import jax
 import jax.numpy as jnp
 import numpy as np
-from ring import algebra
-from ring import base
-from ring import maths
 from tree_utils import PyTree
 from tree_utils import tree_batch
 from vispy import scene
 from vispy.scene import MatrixTransform
 
+from ring import algebra
+from ring import base
+from ring import maths
+
 from . import vispy_visuals
 
 Camera = TypeVar("Camera")
@@ -192,7 +193,7 @@ class Scene(ABC):
 
         # step 3: update visuals
         for i, (visual, geom) in enumerate(zip(self.visuals, self.geoms)):
-            t = jax.tree_map(lambda arr: arr[i], transform_per_visual)
+            t = jax.tree.map(lambda arr: arr[i], transform_per_visual)
             if self._fresh_init:
                 self._init_visual(visual, t, geom)
             else:

ring/sys_composer/inject_sys.py

@@ -2,12 +2,13 @@ from typing import Optional
 
 import jax
 import jax.numpy as jnp
-from ring import base
 from tree_utils import tree_batch
 
+from ring import base
+
 
 def _tree_nan_like(tree, repeats: int):
-    return jax.tree_map(
+    return jax.tree.map(
         lambda arr: jnp.repeat(arr[0:1] * jnp.nan, repeats, axis=0), tree
     )
 

ring/utils/batchsize.py

@@ -39,19 +39,19 @@ def merge_batchsize(
     tree: PyTree, pmap_size: int, vmap_size: int, third_dim_also: bool = False
 ) -> PyTree:
     if third_dim_also:
-        return jax.tree_map(
+        return jax.tree.map(
             lambda arr: arr.reshape(
                 (pmap_size * vmap_size * arr.shape[2],) + arr.shape[3:]
             ),
             tree,
         )
-    return jax.tree_map(
+    return jax.tree.map(
         lambda arr: arr.reshape((pmap_size * vmap_size,) + arr.shape[2:]), tree
     )
 
 
 def expand_batchsize(tree: PyTree, pmap_size: int, vmap_size: int) -> PyTree:
-    return jax.tree_map(
+    return jax.tree.map(
         lambda arr: arr.reshape(
             (
                 pmap_size,

ring/utils/dataloader.py

@@ -4,14 +4,15 @@ from typing import Callable, Optional
 
 import jax
 import numpy as np
-from ring.utils import parse_path
-from ring.utils import pickle_load
 import torch
 from torch.utils.data import DataLoader
 from torch.utils.data import Dataset
 import tqdm
 from tree_utils import PyTree
 
+from ring.utils import parse_path
+from ring.utils import pickle_load
+
 
 def make_generator(
     *paths,
@@ -103,7 +104,7 @@ def pytorch_generator(
     dl_iter = iter(dl)
 
     def to_numpy(tree: PyTree[torch.Tensor]):
-        return jax.tree_map(lambda tensor: tensor.numpy(), tree)
+        return jax.tree.map(lambda tensor: tensor.numpy(), tree)
 
     def generator(_):
         nonlocal dl, dl_iter

ring/utils/dataloader_torch.py

@@ -1,16 +1,25 @@
 import os
+import pickle
 from typing import Any, Optional
 import warnings
 
-import jax
 import numpy as np
 import torch
 from torch.utils.data import DataLoader
 from torch.utils.data import Dataset
+import tree
 from tree_utils import PyTree
 
-from ring.utils import parse_path
-from ring.utils import pickle_load
+from ring.utils.path import parse_path
+
+
+def pickle_load(
+    path,
+):
+    path = parse_path(path, extension="pickle", require_is_file=True)
+    with open(path, "rb") as file:
+        obj = pickle.load(file)
+    return obj
 
 
 class FolderOfFilesDataset(Dataset):
@@ -60,8 +69,8 @@ def dataset_to_generator(
     )
     dl_iter = iter(dl)
 
-    def to_numpy(tree: PyTree[torch.Tensor]):
-        return jax.tree_map(lambda tensor: tensor.numpy(), tree)
+    def to_numpy(data: PyTree[torch.Tensor]):
+        return tree.map_structure(lambda tensor: tensor.numpy(), data)
 
     def generator(_):
         nonlocal dl, dl_iter

ring/utils/hdf5.py

@@ -121,7 +121,7 @@ def _parse_path(
 
 def _tree_concat(trees: list):
     # otherwise scalar-arrays will lead to indexing error
-    trees = jax.tree_map(lambda arr: np.atleast_1d(arr), trees)
+    trees = jax.tree.map(lambda arr: np.atleast_1d(arr), trees)
 
     if len(trees) == 0:
         return trees

ring/utils/normalizer.py

@@ -3,9 +3,10 @@ from typing import Callable, TypeVar
 
 import jax
 import jax.numpy as jnp
-from ring.algorithms.generator import types
 import tree_utils
 
+from ring.algorithms.generator import types
+
 KEY = jax.random.PRNGKey(777)
 KEY_PERMUTATION = jax.random.PRNGKey(888)
 
@@ -37,12 +38,12 @@ def make_normalizer_from_generator(
     # permute 0-th axis, since batchsize of generator might be larger than
     # `approx_with_large_batchsize`, then we would not get a representative
     # subsample otherwise
-    Xs = jax.tree_map(lambda arr: jax.random.permutation(KEY_PERMUTATION, arr), Xs)
+    Xs = jax.tree.map(lambda arr: jax.random.permutation(KEY_PERMUTATION, arr), Xs)
     Xs = tree_utils.tree_slice(Xs, start=0, slice_size=approx_with_large_batchsize)
 
     # obtain statistics
-    mean = jax.tree_map(lambda arr: jnp.mean(arr, axis=(0, 1)), Xs)
-    std = jax.tree_map(lambda arr: jnp.std(arr, axis=(0, 1)), Xs)
+    mean = jax.tree.map(lambda arr: jnp.mean(arr, axis=(0, 1)), Xs)
+    std = jax.tree.map(lambda arr: jnp.std(arr, axis=(0, 1)), Xs)
 
     if verbose:
         print("Mean: ", mean)
@@ -51,6 +52,6 @@ def make_normalizer_from_generator(
     eps = 1e-8
 
     def normalizer(X):
-        return jax.tree_map(lambda a, b, c: (a - b) / (c + eps), X, mean, std)
+        return jax.tree.map(lambda a, b, c: (a - b) / (c + eps), X, mean, std)
 
     return normalizer