imt-ring 1.6.36__py3-none-any.whl → 1.6.38__py3-none-any.whl

{imt_ring-1.6.36.dist-info → imt_ring-1.6.38.dist-info}/METADATA

@@ -1,6 +1,6 @@
-Metadata-Version: 2.1
+Metadata-Version: 2.2
 Name: imt-ring
-Version: 1.6.36
+Version: 1.6.38
 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.36.dist-info → imt_ring-1.6.38.dist-info}/RECORD

@@ -1,25 +1,25 @@
 ring/__init__.py,sha256=H1Rd2uXVkux4Z792XyHIkQ8OpDSZBiPqFwyAFDWDU3E,5260
 ring/algebra.py,sha256=F0GwbP8LQP5qGVkoMUYJmkp9Hn2nKAVIkCVYDEjNjGU,3128
-ring/base.py,sha256=_TgFrggsZfam0VPxvD4J5xp977vgiLnKTlDIJVzik5M,35362
+ring/base.py,sha256=4Yxk6jk-B4UUm_6YYshxmHSHqOg0mhTOxtZP5fFS8nw,35373
 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
 ring/algorithms/_random.py,sha256=UMyv-VPZLcErrKqs0XB83QJjs8GrmoNsv-zRSxGXvnI,14490
 ring/algorithms/dynamics.py,sha256=dpe-F3Yq4sY2dY6DQW3v7TnPLRdxdkePtdbGPQIrijg,10997
 ring/algorithms/jcalc.py,sha256=QafnCKa1mjEl7nL_KuadPJB5ebW31NKnkdcKn2YtSsM,36171
-ring/algorithms/kinematics.py,sha256=DOboHI517Vx0pRJUFZtZPmK_qFaiKiQe-37B-M0aC-c,7422
-ring/algorithms/sensors.py,sha256=0xOzdQIc1kBF0CkoPXWWCx3MmV4SG3wj7knVnnMWq9M,18124
+ring/algorithms/kinematics.py,sha256=IXeTQ-afzeEzLVmQVQ1oTXJxz_lTwvaWlgHeJkhO_8o,7423
+ring/algorithms/sensors.py,sha256=v_TZMyWjffDpPwoyS1fy8X-1i9y1vDf6mk1EmGS2ztc,18251
 ring/algorithms/custom_joints/__init__.py,sha256=3pQ-Is_HBTQDkzESCNg9VfoP8wvseWmooryG8ERnu_A,366
 ring/algorithms/custom_joints/rr_imp_joint.py,sha256=_YJK0p8_0MHFtr1NuGnNZoxTbwaMQyUjYv7EtsPiU3A,2402
 ring/algorithms/custom_joints/rr_joint.py,sha256=jnRtjtOCALMaq2_0bcu2d7qgfQ6etXpoh43MioRaDmY,1000
 ring/algorithms/custom_joints/rsaddle_joint.py,sha256=QoMo6NXdYgA9JygSzBvr0eCdd3qKhUgCrGPNO2Qdxko,1200
-ring/algorithms/custom_joints/suntay.py,sha256=tOEGM304XciHO4pmvxr4faA4xXVO4N2HlPdFmXKbcrw,16726
+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=jGQocoNZ5tkiMazBDCv-jD6FNYwebqn0_RgVFse49pg,16890
-ring/algorithms/generator/batch.py,sha256=P51UnAZl9TUF_eVq58VL1CsmPPStPHhRDdKjUyvu4EA,2652
-ring/algorithms/generator/finalize_fns.py,sha256=nY2RKiLbHriTkdec94lc4UGSZKd0v547MDNn4dr8I3E,10398
+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=XJ_Gd5AkIRh-jBrMfQyMXjVwhx2gCNHznjzFbmAwhZs,5767
+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/io/__init__.py,sha256=1gEJdyDCbldbbm8QeZbLmhzSKmaQ-UqTmQgu4DBH2Z4,328
@@ -47,46 +47,46 @@ ring/io/examples/test_morph_system/four_seg_seg1.xml,sha256=XJvGtEnvedejs_OmCVfQ
 ring/io/examples/test_morph_system/four_seg_seg3.xml,sha256=HktN7_a_Ly3YflWit5W-WncxApWGMORAGnRXyMEqnoA,1265
 ring/io/xml/__init__.py,sha256=-3k6ffvFyc4zm0oTyVz3ez-o3Lb9bPp2sjwSub_K1AA,242
 ring/io/xml/abstract.py,sha256=8Q2ebnUYLmuS9HJAQwDVrDTrRfD5z4G5RAB7MW8Oa60,9742
-ring/io/xml/from_xml.py,sha256=E7JQl_scL5U4LK6mqLMr5qaiZCc6J1fInxD7uwgNCJY,9356
+ring/io/xml/from_xml.py,sha256=CR3OaBxoDuHK8k5N79XziHwS90lCaaw49UGzQirWiIw,9356
 ring/io/xml/test_from_xml.py,sha256=bckVrVVmEhCwujd_OF9FGYnX3zU3BgztpqGxxmd0htM,1562
 ring/io/xml/test_to_xml.py,sha256=NGn4VSiFdwhYN5YTBduWMiY9B5dwtxZhCQAR_PXeqKU,946
 ring/io/xml/to_xml.py,sha256=Wo4iySLw9nM-iVW42AGvMRqjtU2qRc2FD_Zlc7w1IrE,3438
 ring/ml/__init__.py,sha256=nbh48gaswWeY4S4vT1sply_3ROj2DQ7agjoLR4Ho3T8,1517
-ring/ml/base.py,sha256=lfwEZLBDglOSRWChUHoH1kezefhttPV9TMEpNIqsMNw,9972
+ring/ml/base.py,sha256=eEpLuXlhFoJ-cpXoSGjctLaYduQhnSVpbv-FEYftNRs,9972
 ring/ml/callbacks.py,sha256=oCPXl4_Zcw3g0KRgyyUDmdiGxV0phnDVc_t8rEG4Lls,13737
-ring/ml/ml_utils.py,sha256=M--qkXRnhU7tHvgfTHfT9gyY0nhj3zMGEaK0X0drFLs,10915
+ring/ml/ml_utils.py,sha256=hu189AnHcmkhkpEPZZ19O0gWz3T-YKpWQW9buqDTMow,10915
 ring/ml/optimizer.py,sha256=TZF0_LmnewzmGVso-zIQJtpWguUW0fW3HeRpIdG_qoI,4763
-ring/ml/ringnet.py,sha256=mef7jyN2QcApJmQGH3HYZyTV-00q8YpsYOKhW0-ku1k,8973
+ring/ml/ringnet.py,sha256=uybMtLdQV0oTFldlbffrqku_l79y3coJYHwxX8P9QZQ,8973
 ring/ml/rnno_v1.py,sha256=2qE08OIvTJ5PvSxKpYGzGSrvEImWrdAT_qslZ7jP5tA,1372
-ring/ml/train.py,sha256=Da89HxiqXC7xuX2ldpTrJStqKWN-6Vcpml4PPQuihN4,10989
+ring/ml/train.py,sha256=_CtQM3w9L01V5yn23lz0aaIPJN5sOWlL9e7G_9__11c,10989
 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=bSj1_7YL8wZV6cp9oD2CvbkZRSuxVhmPBA2JECxrnUE,8426
-ring/rendering/vispy_render.py,sha256=QmRyA7Hqk3uS1SKjcncwc4_vd1m4yWryW2X0i4jRvCw,10260
+ring/rendering/mujoco_render.py,sha256=HMvZc04I0-lXPBL3hcnBzV2bNiXQAQM7QcHlG_Obmj4,8757
+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
 ring/sim2real/sim2real.py,sha256=B4nqBBnjGXhM-7PfTyxEq44ZidGNghqaq--qdFILX5A,9675
 ring/sys_composer/__init__.py,sha256=5J_JJJIHfTPcpxh0v4FqiOs81V1REPUd7pgiw2nAN5E,193
 ring/sys_composer/delete_sys.py,sha256=cIM9KbyLfg7B9121g7yjzuFbjeNu9cil1dPavAYEgzk,3408
-ring/sys_composer/inject_sys.py,sha256=Mj-q-mUjXKwkg-ol6IQAjf9IJfk7pGhez0_WoTKTgm0,3503
+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/batchsize.py,sha256=FbOii7MDP4oPZd9GJOKehFatfnb6WZ0b9z349iZYs1A,1786
+ring/utils/batchsize.py,sha256=uCj8LG7elbjEUUzuK29Z3I9T8bxJTcsybY3DdGeqhQs,1786
 ring/utils/colab.py,sha256=ZLHwP0jNQUsmZJU4l68a5djULPi6T-jYNNHevjIoMn8,1631
-ring/utils/dataloader.py,sha256=2CcsbUY2AZs8LraS5HTJXlEseuF-1gKmfyBkSsib-tE,3748
-ring/utils/dataloader_torch.py,sha256=bravdBqbkxxcQDieg6OnmArGwGcWpMI3MmNFwTCt0qg,3808
-ring/utils/hdf5.py,sha256=BzXwVypZmEZeHVgeGZ78YYdi10NEQtnPhdrb8dQAXo0,5856
-ring/utils/normalizer.py,sha256=67L2BU1MRsMT4pD41ta3JJMppLN0ozFmnwrmXDtnqrQ,1698
+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=tJaWXLGOTwkxJQj2l23dX97wO3aZYhM2qd7eNuMRs84,6907
+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.36.dist-info/METADATA,sha256=a-uW_s0jWJEBX9kW1q36Br4SNXPK7eGIVhlsyKDWruE,4251
-imt_ring-1.6.36.dist-info/WHEEL,sha256=PZUExdf71Ui_so67QXpySuHtCi3-J3wvF4ORK6k_S8U,91
-imt_ring-1.6.36.dist-info/top_level.txt,sha256=EiT790-lAyi8iwTzJArH3f2k77rwhDn00q-4PlmvDQo,5
-imt_ring-1.6.36.dist-info/RECORD,,
+imt_ring-1.6.38.dist-info/METADATA,sha256=9rN1VzsIlGU8eyABz9-pTxj0OTCFOZRilEEzkB4gyvg,4251
+imt_ring-1.6.38.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
+imt_ring-1.6.38.dist-info/top_level.txt,sha256=EiT790-lAyi8iwTzJArH3f2k77rwhDn00q-4PlmvDQo,5
+imt_ring-1.6.38.dist-info/RECORD,,

{imt_ring-1.6.36.dist-info → imt_ring-1.6.38.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (75.6.0)
+Generator: setuptools (75.8.0)
 Root-Is-Purelib: true
 Tag: py3-none-any
 

ring/algorithms/custom_joints/suntay.py

@@ -184,7 +184,7 @@ def register_suntay(sconfig: SuntayConfig, name: str = "suntay"):
 
         suntay_link_name = _utils_find_suntay_joint(sys)
 
-        params = jax.tree_map(
+        params = jax.tree.map(
             lambda arr: arr[sys.idx_map("l")[suntay_link_name]],
             sys.links.joint_params[name],
         )

ring/algorithms/generator/batch.py

@@ -80,11 +80,11 @@ def generators_eager(
             # 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 = 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]
-            callback([jax.tree_map(lambda a: a[i].copy(), sample) for i in range(size)])
+            callback([jax.tree.map(lambda a: a[i].copy(), sample) for i in range(size)])
 
             # cleanup
             del sample, sample_flat

ring/algorithms/generator/finalize_fns.py

@@ -311,7 +311,7 @@ def _expand_then_flatten(Xy):
 
     X, y = _flatten(X), _flatten(y)
     if not batched:
-        X, y = jax.tree_map(lambda arr: arr[0], (X, y))
+        X, y = jax.tree.map(lambda arr: arr[0], (X, y))
     return X, y
 
 

ring/algorithms/generator/pd_control.py

@@ -86,7 +86,7 @@ def _pd_control(P: jax.Array, D: Optional[jax.Array] = None):
         controller_state: PDControllerState, sys: base.System, state: base.State
     ) -> jax.Array:
         taus = jnp.zeros((sys.qd_size()))
-        q_ref, qd_ref = jax.tree_map(
+        q_ref, qd_ref = jax.tree.map(
             lambda arr: jax.lax.dynamic_index_in_dim(
                 arr, controller_state.i, keepdims=False
             ),

ring/algorithms/kinematics.py

@@ -4,6 +4,7 @@ import jax
 import jax.numpy as jnp
 import jaxopt
 from jaxopt._src.base import Solver
+
 from ring import algebra
 from ring import base
 from ring import maths
@@ -171,7 +172,7 @@ def inverse_kinematics_endeffector(
 
             # find result of best q0 initial value
             best_q_index = jnp.argmin(values)
-            best_q, best_q_value = jax.tree_map(
+            best_q, best_q_value = jax.tree.map(
                 lambda arr: jax.lax.dynamic_index_in_dim(
                     arr, best_q_index, keepdims=False
                 ),

ring/algorithms/sensors.py

@@ -244,7 +244,7 @@ def imu(
         measurements["mag"] = magnetometer(xs.rot, magvec)
 
     if smoothen_degree is not None:
-        measurements = jax.tree_map(
+        measurements = jax.tree.map(
             lambda arr: _moving_average(arr, smoothen_degree),
             measurements,
         )
@@ -257,7 +257,7 @@ def imu(
             delay = half_window
 
     if delay is not None and delay > 0:
-        measurements = jax.tree_map(
+        measurements = jax.tree.map(
             lambda arr: (jnp.pad(arr, ((delay, 0), (0, 0)))[:-delay]), measurements
         )
 
@@ -473,7 +473,7 @@ def _joint_axes_from_sys(sys: base.Transform, N: int) -> dict:
         X[name] = {"joint_axes": joint_axes}
 
     sys.scan(f, "lll", sys.link_names, sys.link_types, sys.links)
-    X = jax.tree_map(lambda arr: jnp.repeat(arr[None], N, axis=0), X)
+    X = jax.tree.map(lambda arr: jnp.repeat(arr[None], N, axis=0), X)
     return X
 
 
@@ -498,12 +498,12 @@ _quasi_physical_sys_str = r"""
 <x_xy>
     <options gravity="0 0 0"/>
     <worldbody>
-        <body name="IMU" joint="p3d" damping="0.1 0.1 0.1" spring_stiff="3 3 3">
-            <geom type="box" mass="0.002" dim="0.01 0.01 0.01"/>
+        <body name="IMU" joint="free" damping="1 1 1 10 10 10" spring_stiff="20 20 20 500 500 500">
+            <geom type="box" mass="1" dim="0.01 0.01 0.01"/>
         </body>
     </worldbody>
 </x_xy>
-"""
+"""  # noqa: E501
 
 
 def _quasi_physical_simulation_beautiful(
@@ -512,12 +512,14 @@ def _quasi_physical_simulation_beautiful(
     sys = io.load_sys_from_str(_quasi_physical_sys_str).replace(dt=dt)
 
     def step_dynamics(state: base.State, x):
-        state = algorithms.step(sys.replace(link_spring_zeropoint=x.pos), state)
+        state = algorithms.step(
+            sys.replace(link_spring_zeropoint=jnp.concatenate((x.rot, x.pos))), state
+        )
         return state, state.q
 
-    state = base.State.create(sys, q=xs.pos[0])
-    _, pos = jax.lax.scan(step_dynamics, state, xs)
-    return xs.replace(pos=pos)
+    state = base.State.create(sys, q=jnp.concatenate((xs.rot[0], xs.pos[0])))
+    _, qs = jax.lax.scan(step_dynamics, state, xs)
+    return xs.replace(rot=qs[:, :4], pos=qs[:, 4:])
 
 
 _constants = {

ring/base.py

@@ -807,7 +807,7 @@ class System(_Base):
         if exists:
             return cls.from_xml(path, seed=seed)
         else:
-            return cls.from_str(path_or_str)
+            return cls.from_str(path_or_str, seed=seed)
 
     def coordinate_vector_to_q(
         self,

ring/io/xml/from_xml.py

@@ -252,7 +252,7 @@ def load_sys_from_str(xml_str: str, seed: int = 1) -> base.System:
 
     # numpy -> jax
     # we load using numpy in order to have float64 precision
-    sys = jax.tree_map(jax.numpy.asarray, sys)
+    sys = jax.tree.map(jax.numpy.asarray, sys)
 
     sys = jcalc._init_joint_params(jax.random.PRNGKey(seed), sys)
 

ring/ml/base.py

@@ -13,13 +13,13 @@ from ring.utils import pickle_save
 def _to_3d(tree):
     if tree is None:
         return None
-    return jax.tree_map(lambda arr: arr[None], tree)
+    return jax.tree.map(lambda arr: arr[None], tree)
 
 
 def _to_2d(tree, i: int = 0):
     if tree is None:
         return None
-    return jax.tree_map(lambda arr: arr[i], tree)
+    return jax.tree.map(lambda arr: arr[i], tree)
 
 
 class AbstractFilter(ABC):

ring/ml/ml_utils.py

@@ -161,7 +161,7 @@ def _flatten_convert_filter_nested_dict(
     metrices: NestedDict, filter_nan_inf: bool = True
 ):
     metrices = _flatten_dict(metrices)
-    metrices = jax.tree_map(_to_float_if_not_string, metrices)
+    metrices = jax.tree.map(_to_float_if_not_string, metrices)
 
     if not filter_nan_inf:
         return metrices
@@ -216,7 +216,7 @@ def save_model_tf(jax_func, path: str, *input, validate: bool = True):
     from jax.experimental import jax2tf
     import tensorflow as tf
 
-    signature = jax.tree_map(
+    signature = jax.tree.map(
         lambda arr: tf.TensorSpec(list(arr.shape), tf.float32), input
     )
 
@@ -241,7 +241,7 @@ def save_model_tf(jax_func, path: str, *input, validate: bool = True):
     if validate:
         output_jax = jax_func(*input)
         output_tf = tf.saved_model.load(path)(*input)
-        jax.tree_map(
+        jax.tree.map(
             lambda a1, a2: np.allclose(a1, a2, atol=1e-5, rtol=1e-5),
             output_jax,
             output_tf,

ring/ml/ringnet.py

@@ -248,7 +248,7 @@ class RING(ml_base.AbstractFilter):
         params, state = self.forward_lam_factory(lam=lam).init(key, X)
 
         if bs is not None:
-            state = jax.tree_map(lambda arr: jnp.repeat(arr[None], bs, axis=0), state)
+            state = jax.tree.map(lambda arr: jnp.repeat(arr[None], bs, axis=0), state)
 
         return params, state
 

ring/ml/train.py

@@ -50,7 +50,7 @@ def _build_step_fn(
         # this vmap maps along batch-axis, not time-axis
         # time-axis is handled by `metric_fn`
         pipe = lambda q, qhat: jnp.mean(jax.vmap(metric_fn)(q, qhat))
-        error_tree = jax.tree_map(pipe, y, yhat)
+        error_tree = jax.tree.map(pipe, y, yhat)
         return jnp.mean(tree_utils.batch_concat(error_tree, 0)), state
 
     @partial(
@@ -274,7 +274,7 @@ def _build_eval_fn(
             ), f"The metric identitifier {metric_name} is not unique"
 
             pipe = lambda q, qhat: reduce_fn(jax.vmap(jax.vmap(metric_fn))(q, qhat))
-            values.update({metric_name: jax.tree_map(pipe, y, yhat)})
+            values.update({metric_name: jax.tree.map(pipe, y, yhat)})
 
         return values
 

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

ring/utils/utils.py

@@ -1,5 +1,6 @@
 from importlib import import_module as _import_module
 import io
+from pathlib import Path
 import pickle
 import random
 from typing import Optional
@@ -152,13 +153,28 @@ def import_lib(
 
 def pickle_save(obj, path, overwrite: bool = False):
     path = parse_path(path, extension="pickle", file_exists_ok=overwrite)
-    with open(path, "wb") as file:
-        pickle.dump(obj, file, protocol=5)
+    try:
+        with open(path, "wb") as file:
+            pickle.dump(obj, file, protocol=5)
+    except OSError as e:
+        print(
+            f"saving with `pickle` throws exception {e}. "
+            + "Attempting to save using `joblib`"
+        )
+        path = parse_path(path, extension="joblib", file_exists_ok=overwrite)
+        import joblib
+
+        joblib.dump(obj, path)
 
 
 def pickle_load(
     path,
 ):
+    if Path(path).suffix == ".joblib":
+        import joblib
+
+        return joblib.load(path)
+
     path = parse_path(path, extension="pickle", require_is_file=True)
     with open(path, "rb") as file:
         obj = pickle.load(file)