unienv 0.0.1b4__py3-none-any.whl → 0.0.1b6__py3-none-any.whl

{unienv-0.0.1b4.dist-info → unienv-0.0.1b6.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: unienv
-Version: 0.0.1b4
+Version: 0.0.1b6
 Summary: Unified robot environment framework supporting multiple tensor and simulation backends
 License-Expression: MIT
 Project-URL: Homepage, https://github.com/UniEnvOrg/UniEnv
@@ -12,9 +12,10 @@ Requires-Python: >=3.10
 Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: numpy
-Requires-Dist: xbarray>=0.0.1a8
+Requires-Dist: xbarray>=0.0.1a13
 Requires-Dist: pillow
 Requires-Dist: cloudpickle
+Requires-Dist: pyvers
 Provides-Extra: dev
 Requires-Dist: pytest; extra == "dev"
 Provides-Extra: gymnasium

{unienv-0.0.1b4.dist-info → unienv-0.0.1b6.dist-info}/RECORD

@@ -1,27 +1,34 @@
-unienv-0.0.1b4.dist-info/licenses/LICENSE,sha256=nkklvEaJUR4QDBygz7tkEe1FMVKV1JSjnGzJNLhdIWM,1091
+unienv-0.0.1b6.dist-info/licenses/LICENSE,sha256=nkklvEaJUR4QDBygz7tkEe1FMVKV1JSjnGzJNLhdIWM,1091
 unienv_data/__init__.py,sha256=zFxbe7aM5JvYXIK0FGnOPwWQJMN-8l_l8prB85CkcA8,95
 unienv_data/base/__init__.py,sha256=w-I8A-z7YYArkHc2ZOVGrfzfThsaDBg7aD7qMFprNM8,186
-unienv_data/base/common.py,sha256=EYOzuYmvsCy1uJftsw6cXeycPIr8P7GWZ3_q4wgoNeo,12879
-unienv_data/base/storage.py,sha256=s99PYEZGa76kf-Enx57kOyVkwjb-tpU-vTHcGc5Dcew,5415
+unienv_data/base/common.py,sha256=A3RtD3Omqk0Qplsc-44ukAEzbQEU22_MkwUlC7l-HHM,13083
+unienv_data/base/storage.py,sha256=afICsO_7Zbm9azV0Jxho_z9F7JM30TUDjJM1NHETDHM,5495
 unienv_data/batches/__init__.py,sha256=Vi92f8ddgFYCqwv7xO2Pi3oJePnioJ4XrJbQVV7eIvk,234
-unienv_data/batches/backend_compat.py,sha256=7Juf7nU2jYHohRzNzmGfqMMpJtFM-3oTzzLu6EbC77E,8168
-unienv_data/batches/combined_batch.py,sha256=aua1H86sa_qWrCtAAp5JIZMGtFiiKFPFkU0y5JoyElM,15325
+unienv_data/batches/backend_compat.py,sha256=tzFG8gTq0yW-J6PLvu--lCGS0lFc0QfelicJ50p_HYc,8207
+unienv_data/batches/combined_batch.py,sha256=pNrbLvU565BUDWO0pZLCnSMygmoGVCLxjC9OkLRKtLA,15330
 unienv_data/batches/framestack_batch.py,sha256=pdURqZeksOlbf21Nhx8kkm0gtFt6rjt2OiNWgZPdFCM,2312
-unienv_data/batches/slicestack_batch.py,sha256=J2EhARcPA-zz6EBnV7OLzm4yyvnZ06vrdUoPD5RkJ-o,16672
+unienv_data/batches/slicestack_batch.py,sha256=Q3-gsJTvMjKTeZAHWNBTGRsws0HctsfMMTw0vylNxvA,16785
 unienv_data/batches/transformations.py,sha256=b4HqX3wZ6TuRgQ2q81Jv43PmeHGmP8cwURK_ULjGNgs,5647
 unienv_data/integrations/pytorch.py,sha256=pW5rXBXagfzwJjM_VGgg8CPXEs3e2fKgg4nY7M3dpOc,2350
 unienv_data/replay_buffer/__init__.py,sha256=uVebYruIYlj8OjTYVi8UYI4gWp3S3XIdgFlHbwO260o,100
-unienv_data/replay_buffer/replay_buffer.py,sha256=nhbC-7aHGIYhcCdmaaDdhB2U9ODAZrbKMq8dP8ffOv0,10344
-unienv_data/replay_buffer/trajectory_replay_buffer.py,sha256=fxV6FIqAHhN8opYs2WjAJMPqNRWD3iIku-4WlaydyG4,20737
+unienv_data/replay_buffer/replay_buffer.py,sha256=8vPma5dL6jDGhI3Oo6IEvNcDYJG9Lb0Xlvxp45tQMEs,14498
+unienv_data/replay_buffer/trajectory_replay_buffer.py,sha256=cqRmzdewFS8IvJcMwxxQgwZf7TvvrViym87OaCOes3Y,24009
 unienv_data/samplers/__init__.py,sha256=e7uunWN3r-g_2fDaMsYMe8cZcF4N-okCxqBPweQnE0s,97
 unienv_data/samplers/multiprocessing_sampler.py,sha256=FEBK8pMTnkpA0xuMkbvlv4aIdVTTubeT8BjL60BJL5o,13254
 unienv_data/samplers/step_sampler.py,sha256=ZCcrx9WbILtaR6izhIP3DhtmFcP7KQBdaYaSZ7vWwRk,3010
-unienv_data/storages/dict_storage.py,sha256=SqCGcGT9Y4l0thdmx23XSxRMzIEIuldA6m8Cd9HrpnA,12588
-unienv_data/storages/flattened.py,sha256=Fu01TjrzvmyNhXEGtC4FiBTb7cqXDtVkErc1QNwLvcI,6704
-unienv_data/storages/hdf5.py,sha256=F_mkrmX6SGT2HamJAyYopBmj_Nf5NzJiyvVN9irtiiM,26260
-unienv_data/storages/pytorch.py,sha256=ftO8cND7PFV0J1B1o2YOWqj4U_pyWsJvWv9lC9A7LJg,6953
-unienv_data/storages/transformation.py,sha256=9BIwrvdruiTRduqC03e5UbSjBT1jLSxLCkNfrsVDP7I,7577
-unienv_data/transformations/image_compress.py,sha256=dINrvmpTWy3sbqruHk0kPZG2XNyJI90ERgErXV7GamE,9131
+unienv_data/storages/_episode_storage.py,sha256=OpZt4P-P6LHrBR4F-tNcCFROLskWaOKWCDfoPV7qz1I,21970
+unienv_data/storages/_list_storage.py,sha256=pH9xZOqXCx65NBRRD-INcP8OP-NWsI-JvdzVsPj9MSg,5225
+unienv_data/storages/backend_compat.py,sha256=BxeMJlC3FI60KLJ7QB5kF-mrGlJ6xi584Dcu4IN4Zrc,10714
+unienv_data/storages/dict_storage.py,sha256=DSqRIgo3m1XtUcLtyjYSqqpi01mr_nJOLg5BCddwPcg,13862
+unienv_data/storages/flattened.py,sha256=Yf1G4D6KE36sESyDMGWKXqhFjz6Idx7N1aEhihmGovA,7055
+unienv_data/storages/hdf5.py,sha256=Jnls1rs7nlOOp9msmAfhuZp80OZd8S2Llls176EOUc4,27096
+unienv_data/storages/image_storage.py,sha256=4J1ZiGFHbGLHmReMztImJoDcRmiB_llD2wbMB3rdvOQ,5137
+unienv_data/storages/npz_storage.py,sha256=IP2DXbUs_ySzILne3s3hq3gwHiy9tfpWz6HcNciA8DU,4868
+unienv_data/storages/pytorch.py,sha256=bf3ys6eBlMvjyPK4XE-itENjEWq5Vm60qNwBNqJIZqg,7345
+unienv_data/storages/transformation.py,sha256=-9_jPZNpx6RXY_ojv_1UCSTa4Z9apI9V9jit8nG93oM,8133
+unienv_data/storages/video_storage.py,sha256=2vcNlghhDZWWzAdf9t0VeCMZrv-x_rYkYaCw8XV8AJA,13331
+unienv_data/third_party/tensordict/memmap_tensor.py,sha256=J6SkFf-FDy43XuaHLgbvDsHt6v2vYfuhRyeoV02P8vw,42589
+unienv_data/transformations/image_compress.py,sha256=f8JTY4DJEXaiu5lO77T4ROV950rh_bOZBchOF-O0tx8,13130
 unienv_interface/__init__.py,sha256=pAWqfm4l7NAssuyXCugIjekSIh05aBbOjNhwsNXcJbE,100
 unienv_interface/backends/__init__.py,sha256=L7CFwCChHVL-2Dpz34pTGC37WgodfJEeDQwXscyM7FM,198
 unienv_interface/backends/base.py,sha256=1_hji1qwNAhcEtFQdAuzaNey9g5bWYj38t1sQxjnggc,132
@@ -36,38 +43,42 @@ unienv_interface/env_base/funcenv_wrapper.py,sha256=chw1iJ1RhAFMv4JAk67cttJvI9ag
 unienv_interface/env_base/vec_env.py,sha256=bcv6NdOxt0Xp1fRMXqzFtmVw6LQ-pDj_Jvj-qaW6otQ,16116
 unienv_interface/env_base/wrapper.py,sha256=7hf4Rr2wouS0igPoahhvb2tzYY3bCaWL0NlgwpYZwQs,9734
 unienv_interface/func_wrapper/__init__.py,sha256=6BPF8O25WkIBpODVTwnOE9HGSm3KRKX6iPwFGWESlxA,123
-unienv_interface/func_wrapper/frame_stack.py,sha256=52CqAHDqwgHtOwMwxzB3Syup9kA19zdlvXCH4mI7MNU,6819
+unienv_interface/func_wrapper/frame_stack.py,sha256=wuGsrluoz60FTczRuo8sHPfpl_Yl4GVTRBb2QDzYPrA,6825
 unienv_interface/func_wrapper/transformation.py,sha256=7mdzcpjLjqtpbtXoqbkGtTMPQxoMmMsqzDWHcZLbrhs,5939
 unienv_interface/space/__init__.py,sha256=6-wLoD9mKDAfz7IuQs_Rn9DMDfDwTZ0tEhQ924libpg,99
 unienv_interface/space/space.py,sha256=mFlCcDvMgEPTXlwo_iwBlm6Eg4Bn2rrecgsfIVstdq0,4067
 unienv_interface/space/space_utils/__init__.py,sha256=GAsPoZC8YNabx3Gw5m2o4zsnG8zmA3mcuM9_lNKhiGo,121
-unienv_interface/space/space_utils/batch_utils.py,sha256=qXK7kERPXKGIYozz7lpjzVz56S9GkH6ZASfIRzCYXHY,36993
+unienv_interface/space/space_utils/batch_utils.py,sha256=hD4ItBp2WQzIQR5u0Zkw0FQQfOeg6ZPRi18Johmcc40,37150
 unienv_interface/space/space_utils/construct_utils.py,sha256=Y4RpV9obY8XQ85O3r_NC1HrBk-Nm941ffRNXNL7nHgA,8323
-unienv_interface/space/space_utils/flatten_utils.py,sha256=kkHkjrsk43NDbg3Q5VAhVoIXStuRayYFO-7knsDzx4A,12289
+unienv_interface/space/space_utils/flatten_utils.py,sha256=6ObJgVq4yhOq_7N5E5pQZS6WmmeKu-MyRFJ_x-gqmNg,12607
 unienv_interface/space/space_utils/gym_utils.py,sha256=nH8EKruOKCXNrIMPUd9F4XGKCfFkhxsTmx4I1BeSgn0,15079
 unienv_interface/space/space_utils/serialization_utils.py,sha256=LWYSFN7E6tEFe8ULWm42LkFUxP_0dfTGkCcx0yl4Y8s,9530
 unienv_interface/space/spaces/__init__.py,sha256=Jap768TlwHFDDiTzHZ0qaHEFEVC1cKA2QzLlSZVQnjI,535
 unienv_interface/space/spaces/batched.py,sha256=RA8aLUSS14zBSCTm_ud18TTa-ntbIZ074xwJ0xls1Jk,3691
 unienv_interface/space/spaces/binary.py,sha256=0iQUbO37dhkznVpjhsJdwlD-KdWgCEx2H7KrybuZ_PM,3570
 unienv_interface/space/spaces/box.py,sha256=NCmileEZOKz-L3WNzZ-uwydrRFsIMdNZBwTn1vWgeI0,13316
-unienv_interface/space/spaces/dict.py,sha256=G5_iYC1Bj5DqeJ7aFlq6eRJbnpATbIRIyRu1jF_UUvk,7022
+unienv_interface/space/spaces/dict.py,sha256=NggllKi0smoz2bL3yrfBM5FJGBNRWZ05xXaNEqY1QKs,7234
 unienv_interface/space/spaces/dynamic_box.py,sha256=HvMNgzfYwIVc5VVgCtq-8lQbNI1V1dZMI-w60AwYss4,19591
 unienv_interface/space/spaces/graph.py,sha256=KocRFLtYP5VWYpwbP6HybXH5R4jTIYJdNePKb6vhnYE,15163
 unienv_interface/space/spaces/text.py,sha256=ePGGJdiD3q-BAX6IHLO7HMe0OH4VrzF043K02eb0zXI,4443
-unienv_interface/space/spaces/tuple.py,sha256=rgZQz3EB3CLbIk9UlHBQbM6w9gssbA1izm-Qq-_Chqs,4267
+unienv_interface/space/spaces/tuple.py,sha256=mmJab6kl5VtQStyn754pmk0RLPSQW06Mu15Hp3Qad80,4287
 unienv_interface/space/spaces/union.py,sha256=Qisd-DdmPcGRmdhZFGiQw8_AOjYWqkuQ4Hwd-I8tdSI,4375
-unienv_interface/transformations/__init__.py,sha256=g19uGnDHMywvDAXRaqFgoWAF1vCPrbJENEpaEgtIrOw,353
-unienv_interface/transformations/batch_and_unbatch.py,sha256=ELCnNtwmgA5wpTBJZasfNSHmtf4vzydzLPmO6IGbT9o,1164
-unienv_interface/transformations/chained_transform.py,sha256=TDnUvxUKK6bXGc_sfr6ZCvvVWw7P5KX2sA9i7i2lx14,2075
-unienv_interface/transformations/dict_transform.py,sha256=ynrJrloVUix2I27Ir1mL86crT0vY5DvpiBAVxPBJup4,5357
+unienv_interface/transformations/__init__.py,sha256=zf8NbY-HW4EgHri9PxpuelEvBpFwUtDEcJiXXhFSDNQ,435
+unienv_interface/transformations/batch_and_unbatch.py,sha256=DK-0VNFvoB9FD1qPtq9SNphQnuzRYJrm_L7VuPiLaMY,2145
+unienv_interface/transformations/chained_transform.py,sha256=_6E1g_8u-WAxKd-f2sHJwKQk9HTIRnulyXwHUwJP12I,2203
+unienv_interface/transformations/crop.py,sha256=sigcQcLklp3P6b6KQfP-Ja3OV1CWeusCLNKMvNNdACQ,3107
+unienv_interface/transformations/dict_transform.py,sha256=GhFSN9t3mL3gvoD_GH-np68Fo4m78YnSyHbUHeyzKcw,5540
 unienv_interface/transformations/filter_dict.py,sha256=DzR-hgHoHJObTipxwB2UrKVlTxbfIrJohaOgqdAICLY,5871
-unienv_interface/transformations/rescale.py,sha256=fM5ukWUvNvPeDO48_PRU0KyyvGhBIDxaN9XZyQ1VaQQ,4364
+unienv_interface/transformations/identity.py,sha256=biW3caBis6ixlOJQk2RJ-7OzP16n0yhpIuqvd7e7Ack,549
+unienv_interface/transformations/image_resize.py,sha256=QyPnpMvdx3IvQyW5_iRq7LMnQQuq7XpOv3x6qQHuNeI,4454
+unienv_interface/transformations/iter_transform.py,sha256=lK7fopeiZJrO0WUXFoUmAOhmYkdHXDnChsQ9TJGV8hU,3688
+unienv_interface/transformations/rescale.py,sha256=85PAq5ta9KelxMaL6RIJXBFxOmRbZjsGlMJiElCW9wI,5329
 unienv_interface/transformations/transformation.py,sha256=u4_9H1tvophhgG0p0F3xfkMMsRuaKY2TQmVeGoeQsJ0,1652
 unienv_interface/utils/control_util.py,sha256=lY_1EknglY3cNekWX9rYWt0ZUglaPMtIt4M5D9y0WfE,2351
-unienv_interface/utils/data_queue.py,sha256=UZiuQDOn39DB9Heu6xinrwuzAL3X8jHlDkFoSC5Phtc,5707
+unienv_interface/utils/framestack_queue.py,sha256=UZiuQDOn39DB9Heu6xinrwuzAL3X8jHlDkFoSC5Phtc,5707
 unienv_interface/utils/seed_util.py,sha256=Up3nBXj7L8w-S9W5Q1U2d9accMhMf0TmHPaN6JXDVWs,677
 unienv_interface/utils/stateclass.py,sha256=xjzicPGX1UuI7q3ZAxhBCCoouKfNtLywUzQtLaT0yS4,1390
-unienv_interface/utils/symbol_util.py,sha256=NAERK-D_2MaTg2eYW-L75tbzPQN5YJIiKtM9zuQ89Sw,383
+unienv_interface/utils/symbol_util.py,sha256=EKC5cVyuXaP5n68-bSbk1A3jCCJCrX90BF7c8mFQYrU,562
 unienv_interface/utils/vec_util.py,sha256=EIK680ReCl_rr-qKP8co5hwz8Dx-gks8SHf-CLOZSOA,373
 unienv_interface/world/__init__.py,sha256=aGuYTz8XFzW32RGkdi2b2LJ1sa0kgFrQyOR3JXDEwLQ,230
 unienv_interface/world/combined_funcnode.py,sha256=O9qWxhtMJkDVtWuGyaeEj3nKMgIyRAPqF9-5LU6yna8,10853
@@ -78,16 +89,16 @@ unienv_interface/world/node.py,sha256=EAvHnx0u7IudmWQDbAUIRVEqB4kh2Xsm1aXdS3Celo
 unienv_interface/world/world.py,sha256=Kl7wbNbs2YR3CjFrCLFhDB3DQUAWM6LjBwSADQtBTII,5740
 unienv_interface/wrapper/__init__.py,sha256=ZNqr-WjVRqgvIxkLkeABxpYZ6tRgJNZOzmluDeJ6W_w,614
 unienv_interface/wrapper/action_rescale.py,sha256=rTJlEHvWSuwGVX83cjfLWvszBk7B2iExX_K37vH8Wic,1231
-unienv_interface/wrapper/backend_compat.py,sha256=T6hosgu2hrZvg3xtnyELmR6Exlz-ztqdj9vdyiz7bhI,7081
+unienv_interface/wrapper/backend_compat.py,sha256=amLAITi1qLylQ45BkpvmwMXSkG-J9YEu1JPjCrBT5I8,7120
 unienv_interface/wrapper/batch_and_unbatch.py,sha256=HpmnppgOKmshNlfmJYkGQYtEU7_U7q3mEdY5n4UaqEY,3457
 unienv_interface/wrapper/control_frequency_limit.py,sha256=B0E2aUbaUr2p2yIN6wT3q4rAbPYsVroioqma2qKMoC0,2322
 unienv_interface/wrapper/flatten.py,sha256=NWA5xne5j_L34oq_wT85wGvp6iHwdCSeGsk1DMugvRw,5837
-unienv_interface/wrapper/frame_stack.py,sha256=lZZh_T_AmxsRWeYSLsTU321lVgIt12MX1eWl_yRNlWg,6002
+unienv_interface/wrapper/frame_stack.py,sha256=07rt8SuUQmniu0HRAzAuSrW9K1ri_87UxxsF-WIUzbI,6008
 unienv_interface/wrapper/gym_compat.py,sha256=JhLxDsO1NsJnKzKhO0MqMw9i5_1FLxoxKilWaQQyBkw,9789
 unienv_interface/wrapper/time_limit.py,sha256=VRvB00BK7deI2QtdGatqwDWmPgjgjg1E7MTvEyaW5rg,2904
 unienv_interface/wrapper/transformation.py,sha256=pQ-_YVU8WWDqSk2sONUUgQY1iigOD092KNcp1DYxoxk,10043
 unienv_interface/wrapper/video_record.py,sha256=y_nJRYgo1SeLeO_Ymg9xbbGPKm48AbU3BxZK2wd0gzk,8679
-unienv-0.0.1b4.dist-info/METADATA,sha256=R_70XnKo1K6ObRxMmSlW1W_lxfD_rGR6txa3wBHGPOM,3033
-unienv-0.0.1b4.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-unienv-0.0.1b4.dist-info/top_level.txt,sha256=wfcJ5_DruUtOEUZjEyfadaKn7B90hWqz2aw-eM3wX5g,29
-unienv-0.0.1b4.dist-info/RECORD,,
+unienv-0.0.1b6.dist-info/METADATA,sha256=i9zJN3WnVG3gmSSZnKI8SROtMN3U5HGUGRD4p1fu3Ho,3056
+unienv-0.0.1b6.dist-info/WHEEL,sha256=qELbo2s1Yzl39ZmrAibXA2jjPLUYfnVhUNTlyF1rq0Y,92
+unienv-0.0.1b6.dist-info/top_level.txt,sha256=wfcJ5_DruUtOEUZjEyfadaKn7B90hWqz2aw-eM3wX5g,29
+unienv-0.0.1b6.dist-info/RECORD,,

{unienv-0.0.1b4.dist-info → unienv-0.0.1b6.dist-info}/WHEEL

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

unienv_data/base/common.py

@@ -7,6 +7,7 @@ from unienv_interface.env_base.env import ContextType, ObsType, ActType
 from unienv_interface.space import Space, BoxSpace, DictSpace
 import dataclasses
 
+from functools import cached_property
 from unienv_interface.space.space_utils import batch_utils as space_batch_utils, flatten_utils as space_flatten_utils
 
 __all__ = [
@@ -46,13 +47,18 @@ class BatchBase(abc.ABC, Generic[BatchT, BArrayType, BDeviceType, BDtypeType, BR
     ):
         self.single_space = single_space
         self.single_metadata_space = single_metadata_space
-        self._batched_space : Space[BatchT, BDeviceType, BDtypeType, BRNGType] = space_batch_utils.batch_space(single_space, 1)
-        if single_metadata_space is not None:
-            self._batched_metadata_space : DictSpace[
-                BDeviceType, BDtypeType, BRNGType
-            ] = space_batch_utils.batch_space(single_metadata_space, 1)
+
+    # For backwards compatibility
+    @cached_property
+    def _batched_space(self) -> Space[BatchT, BDeviceType, BDtypeType, BRNGType]:
+        return space_batch_utils.batch_space(self.single_space, 1)
+    
+    @cached_property
+    def _batched_metadata_space(self) -> Optional[DictSpace[BDeviceType, BDtypeType, BRNGType]]:
+        if self.single_metadata_space is not None:
+            return space_batch_utils.batch_space(self.single_metadata_space, 1)
         else:
-            self._batched_metadata_space = None
+            return None
 
     @property
     def backend(self) -> ComputeBackend[BArrayType, BDeviceType, BDtypeType, BRNGType]:
@@ -146,7 +152,7 @@ class BatchBase(abc.ABC, Generic[BatchT, BArrayType, BDeviceType, BDtypeType, BR
         if tqdm:
             from tqdm import tqdm
             iterable_start = tqdm(iterable_start, desc="Extending Batch")
-        for start_idx in range(0, n_total, chunk_size):
+        for start_idx in iterable_start:
             end_idx = min(start_idx + chunk_size, n_total)
             data_chunk = other.get_at(slice(start_idx, end_idx))
             self.extend(data_chunk)
@@ -183,15 +189,24 @@ class BatchSampler(
     ) -> None:
         super().__init__(single_space=single_space, single_metadata_space=single_metadata_space)
         self.batch_size = batch_size
-        self._batched_space : Space[SamplerBatchT, SamplerDeviceType, SamplerDtypeType, SamplerRNGType] = space_batch_utils.batch_space(self.single_space, batch_size)
-        self._batched_metadata_space : Optional[DictSpace[SamplerDeviceType, SamplerDtypeType, SamplerRNGType]] = space_batch_utils.batch_space(self.single_metadata_space, batch_size) if self.single_metadata_space is not None else None
-
+    
     def manual_seed(self, seed : int) -> None:
         if self.rng is not None:
             self.rng = self.backend.random.random_number_generator(seed, device=self.device)
         if self.data_rng is not None:
             self.data_rng = self.backend.random.random_number_generator(seed, device=self.data.device)
 
+    @cached_property
+    def _batched_space(self) -> Space[BatchT, BDeviceType, BDtypeType, BRNGType]:
+        return space_batch_utils.batch_space(self.single_space, self.batch_size)
+    
+    @cached_property
+    def _batched_metadata_space(self) -> Optional[DictSpace[BDeviceType, BDtypeType, BRNGType]]:
+        if self.single_metadata_space is not None:
+            return space_batch_utils.batch_space(self.single_metadata_space, self.batch_size)
+        else:
+            return None
+    
     @property
     def sampled_space(self) -> Space[SamplerBatchT, SamplerDeviceType, SamplerDtypeType, SamplerRNGType]:
         return self._batched_space

unienv_data/base/storage.py

@@ -20,6 +20,7 @@ class SpaceStorage(abc.ABC, Generic[BatchT, BArrayType, BDeviceType, BDtypeType,
         *args,
         capacity : Optional[int],
         cache_path : Optional[Union[str, os.PathLike]] = None,
+        multiprocessing : bool = False,
         **kwargs
     ) -> "SpaceStorage[BatchT, BArrayType, BDeviceType, BDtypeType, BRNGType]":
         raise NotImplementedError
@@ -32,6 +33,7 @@ class SpaceStorage(abc.ABC, Generic[BatchT, BArrayType, BDeviceType, BDtypeType,
         *,
         capacity : Optional[int] = None,
         read_only : bool = True,
+        multiprocessing : bool = False,
         **kwargs
     ) -> "SpaceStorage[BatchT, BArrayType, BDeviceType, BDtypeType, BRNGType]":
         raise NotImplementedError

unienv_data/batches/backend_compat.py

@@ -36,7 +36,7 @@ def data_to(
             key: data_to(value, source_backend, target_backend, target_device)
             for key, value in data.items()
         }
-    elif isinstance(data, Sequence):
+    elif isinstance(data, Sequence) and not isinstance(data, (str, bytes)):
         data = [
             data_to(value, source_backend, target_backend, target_device)
             for value in data

unienv_data/batches/combined_batch.py

@@ -94,7 +94,7 @@ class CombinedBatch(BatchBase[
         batch_index = int(self.backend.sum(
             idx >= self.index_caches[:, 0]
         ) - 1)
-        return batch_index, idx - self.index_caches[batch_index, 0]
+        return batch_index, idx - int(self.index_caches[batch_index, 0])
 
     def _convert_index(self, idx : Union[IndexableType, BArrayType]) -> Tuple[
         int, 

unienv_data/batches/slicestack_batch.py

@@ -33,6 +33,7 @@ class SliceStackedBatch(BatchBase[
         fill_invalid_data : bool = True,
         stack_metadata : bool = False,
     ):
+        assert batch.backend.dtype_is_real_integer(fixed_offset.dtype), "Fixed offset must be an integer tensor"
         assert len(fixed_offset.shape) == 1, "Fixed offset must be a 1D tensor"
         assert fixed_offset.shape[0] > 0, "Fixed offset must have a positive length"
         assert batch.backend.any(fixed_offset == 0), "There should be at least one zero in the fixed offset"

unienv_data/replay_buffer/replay_buffer.py

@@ -1,6 +1,10 @@
 import abc
 import os
 import dataclasses
+import multiprocessing as mp
+import ctypes
+from contextlib import nullcontext
+
 from typing import Generic, TypeVar, Optional, Any, Dict, Union, Tuple, Sequence, Callable, Type
 from unienv_interface.backends import ComputeBackend, BArrayType, BDeviceType, BDtypeType, BRNGType
 
@@ -51,7 +55,6 @@ def index_with_offset(
         return data_index
 
 class ReplayBuffer(BatchBase[BatchT, BArrayType, BDeviceType, BDtypeType, BRNGType]):
-    is_mutable = True
     # =========== Class Attributes ==========
     @staticmethod
     def create(
@@ -60,6 +63,7 @@ class ReplayBuffer(BatchBase[BatchT, BArrayType, BDeviceType, BDtypeType, BRNGTy
         *args,
         cache_path : Optional[Union[str, os.PathLike]] = None,
         capacity : Optional[int] = None,
+        multiprocessing : bool = False,
         **kwargs
     ) -> "ReplayBuffer[BatchT, BArrayType, BDeviceType, BDtypeType, BRNGType]":
         storage_path_relative = "storage" + (storage_cls.single_file_ext or "")
@@ -70,6 +74,7 @@ class ReplayBuffer(BatchBase[BatchT, BArrayType, BDeviceType, BDtypeType, BRNGTy
             *args,
             cache_path=None if cache_path is None else os.path.join(cache_path, storage_path_relative),
             capacity=capacity,
+            multiprocessing=multiprocessing,
             **kwargs
         )
         return ReplayBuffer(
@@ -77,7 +82,8 @@ class ReplayBuffer(BatchBase[BatchT, BArrayType, BDeviceType, BDtypeType, BRNGTy
             storage_path_relative,
             0,
             0,
-            cache_path=cache_path
+            cache_path=cache_path,
+            multiprocessing=multiprocessing
         )
     
     @staticmethod
@@ -89,6 +95,48 @@ class ReplayBuffer(BatchBase[BatchT, BArrayType, BDeviceType, BDtypeType, BRNGTy
                 metadata = json.load(f)
             return metadata.get('type', None) == __class__.__name__
         return False
+    
+    @staticmethod
+    def get_length_from_path(
+        path : Union[str, os.PathLike]
+    ) -> Optional[int]:
+        if os.path.exists(os.path.join(path, "metadata.json")):
+            with open(os.path.join(path, "metadata.json"), "r") as f:
+                metadata = json.load(f)
+            if metadata.get('type', None) != __class__.__name__:
+                return None
+            return int(metadata["count"])
+        return None
+    
+    @staticmethod
+    def get_capacity_from_path(
+        path : Union[str, os.PathLike]
+    ) -> Optional[int]:
+        if os.path.exists(os.path.join(path, "metadata.json")):
+            with open(os.path.join(path, "metadata.json"), "r") as f:
+                metadata = json.load(f)
+            if metadata.get('type', None) != __class__.__name__:
+                return None
+            return int(metadata["capacity"])
+        return None
+    
+    @staticmethod
+    def get_space_from_path(
+        path : Union[str, os.PathLike],
+        *,
+        backend: ComputeBackend[BArrayType, BDeviceType, BDtypeType, BRNGType],
+        device: Optional[BDeviceType] = None,
+    ) -> Optional[Space[BatchT, BDeviceType, BDtypeType, BRNGType]]:
+        if os.path.exists(os.path.join(path, "metadata.json")):
+            with open(os.path.join(path, "metadata.json"), "r") as f:
+                metadata = json.load(f)
+            if metadata.get('type', None) != __class__.__name__:
+                return None
+            single_instance_space = bsu.json_to_space(
+                metadata["single_instance_space"], backend, device
+            )
+            return single_instance_space
+        return None
 
     @staticmethod
     def load_from(
@@ -97,6 +145,7 @@ class ReplayBuffer(BatchBase[BatchT, BArrayType, BDeviceType, BDtypeType, BRNGTy
         backend: ComputeBackend[BArrayType, BDeviceType, BDtypeType, BRNGType],
         device: Optional[BDeviceType] = None,
         read_only : bool = True,
+        multiprocessing : bool = False,
         **storage_kwargs
     ) -> "ReplayBuffer[BatchT, BArrayType, BDeviceType, BDtypeType, BRNGType]":
         with open(os.path.join(path, "metadata.json"), "r") as f:
@@ -118,52 +167,103 @@ class ReplayBuffer(BatchBase[BatchT, BArrayType, BDeviceType, BDtypeType, BRNGTy
             single_instance_space,
             capacity=capacity,
             read_only=read_only,
+            multiprocessing=multiprocessing,
             **storage_kwargs
         )
-        return ReplayBuffer(storage, metadata["storage_path_relative"], count, offset, cache_path=path)
+        return ReplayBuffer(
+            storage,
+            metadata["storage_path_relative"],
+            count,
+            offset,
+            cache_path=path,
+            multiprocessing=multiprocessing
+        )
 
     # =========== Instance Attributes and Methods ==========
     def dumps(self, path : Union[str, os.PathLike]):
-        os.makedirs(path, exist_ok=True)
-        storage_path = os.path.join(path, self.storage_path_relative)
-        self.storage.dumps(storage_path)
-        metadata = {
-            "type": __class__.__name__,
-            "count": self.count,
-            "offset": self.offset,
-            "capacity": self.storage.capacity,
-            "storage_cls": get_full_class_name(type(self.storage)),
-            "storage_path_relative": self.storage_path_relative,
-            "single_instance_space": bsu.space_to_json(self.storage.single_instance_space),
-        }
-        with open(os.path.join(path, "metadata.json"), "w") as f:
-            json.dump(metadata, f)
+        with self._lock_scope():
+            os.makedirs(path, exist_ok=True)
+            storage_path = os.path.join(path, self.storage_path_relative)
+            self.storage.dumps(storage_path)
+            metadata = {
+                "type": __class__.__name__,
+                "count": self.count,
+                "offset": self.offset,
+                "capacity": self.storage.capacity,
+                "storage_cls": get_full_class_name(type(self.storage)),
+                "storage_path_relative": self.storage_path_relative,
+                "single_instance_space": bsu.space_to_json(self.storage.single_instance_space),
+            }
+            with open(os.path.join(path, "metadata.json"), "w") as f:
+                json.dump(metadata, f)
 
     def __init__(
         self,
         storage : SpaceStorage[BatchT, BArrayType, BDeviceType, BDtypeType, BRNGType],
-        storage_path_relative : Union[str, os.PathLike],
+        storage_path_relative : str,
         count : int = 0,
         offset : int = 0,
         cache_path : Optional[Union[str, os.PathLike]] = None,
+        multiprocessing : bool = False,
     ):
         self.storage = storage
-        self.count = count
-        self.offset = offset
-        self.storage_path_relative = storage_path_relative
+        self._storage_path_relative = storage_path_relative
         self._cache_path = cache_path
+        self._multiprocessing = multiprocessing
+        if multiprocessing and storage.is_mutable:
+            assert storage.is_multiprocessing_safe, "Storage is not multiprocessing safe"
+            self._lock = mp.RLock()
+            self._count_value = mp.Value(ctypes.c_long, int(count))
+            self._offset_value = mp.Value(ctypes.c_long, int(offset))
+        else:
+            self._lock = None
+            self._count_value = int(count)
+            self._offset_value = int(offset)
+        
         super().__init__(
             storage.single_instance_space,
             None
         )
 
+    def _lock_scope(self):
+        if self._lock is not None:
+            return self._lock
+        else:
+            return nullcontext()
+
     @property
     def cache_path(self) -> Optional[Union[str, os.PathLike]]:
         return self._cache_path
 
+    @property
+    def storage_path_relative(self) -> str:
+        return self._storage_path_relative
+
     def __len__(self) -> int:
         return self.count
+
+    @property
+    def count(self) -> int:
+        return self._count_value.value if not isinstance(self._count_value, int) else self._count_value
+
+    @count.setter
+    def count(self, value: int) -> None:
+        if not isinstance(self._count_value, int):
+            self._count_value.value = int(value)
+        else:
+            self._count_value = int(value)
     
+    @property
+    def offset(self) -> int:
+        return self._offset_value.value if not isinstance(self._offset_value, int) else self._offset_value
+
+    @offset.setter
+    def offset(self, value: int) -> None:
+        if not isinstance(self._offset_value, int):
+            self._offset_value.value = int(value)
+        else:
+            self._offset_value = int(value)
+
     @property
     def capacity(self) -> Optional[int]:
         return self.storage.capacity
@@ -176,12 +276,21 @@ class ReplayBuffer(BatchBase[BatchT, BArrayType, BDeviceType, BDtypeType, BRNGTy
     def device(self) -> Optional[BDeviceType]:
         return self.storage.device
 
+    @property
+    def is_mutable(self) -> bool:
+        return self.storage.is_mutable
+    
+    @property
+    def is_multiprocessing_safe(self) -> bool:
+        return self._multiprocessing
+
     def get_flattened_at(self, idx):
         return self.get_flattened_at_with_metadata(idx)[0]
 
     def get_flattened_at_with_metadata(self, idx: Union[IndexableType, BArrayType]) -> BArrayType:
         if hasattr(self.storage, "get_flattened"):
-            data = self.storage.get_flattened(idx)
+            with self._lock_scope():
+                data = self.storage.get_flattened(idx)
             return data, None
 
         data, metadata = self.get_at_with_metadata(idx)
@@ -195,19 +304,21 @@ class ReplayBuffer(BatchBase[BatchT, BArrayType, BDeviceType, BDtypeType, BRNGTy
         return self.get_at_with_metadata(idx)[0]
 
     def get_at_with_metadata(self, idx):
-        data_index = index_with_offset(
-            self.backend,
-            idx,
-            self.count,
-            self.offset,
-            self.device
-        )
-        data = self.storage.get(data_index)
+        with self._lock_scope():
+            data_index = index_with_offset(
+                self.backend,
+                idx,
+                self.count,
+                self.offset,
+                self.device
+            )
+            data = self.storage.get(data_index)
         return data, None
     
     def set_flattened_at(self, idx: Union[IndexableType, BArrayType], value: BArrayType) -> None:
         if hasattr(self.storage, "set_flattened"):
-            self.storage.set_flattened(idx, value)
+            with self._lock_scope():
+                self.storage.set_flattened(idx, value)
             return
 
         if isinstance(idx, int):
@@ -217,13 +328,14 @@ class ReplayBuffer(BatchBase[BatchT, BArrayType, BDeviceType, BDtypeType, BRNGTy
         self.set_at(idx, value)
 
     def set_at(self, idx, value):
-        self.storage.set(index_with_offset(
-            self.backend,
-            idx,
-            self.count,
-            self.offset,
-            self.device
-        ), value)
+        with self._lock_scope():
+            self.storage.set(index_with_offset(
+                self.backend,
+                idx,
+                self.count,
+                self.offset,
+                self.device
+            ), value)
 
     def extend_flattened(
         self,
@@ -233,35 +345,37 @@ class ReplayBuffer(BatchBase[BatchT, BArrayType, BDeviceType, BDtypeType, BRNGTy
         self.extend(unflattened_data)
         
     def extend(self, value):
-        B = sbu.batch_size_data(value)
-        if B == 0:
-            return
-        if self.capacity is None:
-            assert self.offset == 0, "Offset must be 0 when capacity is None"
-            self.storage.extend_length(B)
-            self.storage.set(slice(self.count, self.count + B), value)
-            self.count += B
-            return
-        
-        # We have a fixed capacity, only keep the last `capacity` elements
-        if B >= self.capacity:
-            self.storage.set(Ellipsis, sbu.get_at(self._batched_space, value, slice(-self.capacity, None)))
-            self.count = self.capacity
-            self.offset = 0
-            return
-        
-        # Otherwise, perform round-robin writes
-        indexes = (self.backend.arange(B, device=self.device) + self.offset + self.count) % self.capacity
-        self.storage.set(indexes, value)
-        outflow = max(0, self.count + B - self.capacity)
-        if outflow > 0:
-            self.offset = (self.offset + outflow) % self.capacity
-        self.count = min(self.count + B, self.capacity)
+        with self._lock_scope():
+            B = sbu.batch_size_data(value)
+            if B == 0:
+                return
+            if self.capacity is None:
+                assert self.offset == 0, "Offset must be 0 when capacity is None"
+                self.storage.extend_length(B)
+                self.storage.set(slice(self.count, self.count + B), value)
+                self.count += B
+                return
+            
+            # We have a fixed capacity, only keep the last `capacity` elements
+            if B >= self.capacity:
+                self.storage.set(Ellipsis, sbu.get_at(self._batched_space, value, slice(-self.capacity, None)))
+                self.count = self.capacity
+                self.offset = 0
+                return
+            
+            # Otherwise, perform round-robin writes
+            indexes = (self.backend.arange(B, device=self.device) + self.offset + self.count) % self.capacity
+            self.storage.set(indexes, value)
+            outflow = max(0, self.count + B - self.capacity)
+            if outflow > 0:
+                self.offset = (self.offset + outflow) % self.capacity
+            self.count = min(self.count + B, self.capacity)
 
     def clear(self):
-        self.count = 0
-        self.offset = 0
-        self.storage.clear()
+        with self._lock_scope():
+            self.count = 0
+            self.offset = 0
+            self.storage.clear()
 
     def close(self) -> None:
         self.storage.close()