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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: unienv
-Version: 0.0.1b4
+Version: 0.0.1b5
 Summary: Unified robot environment framework supporting multiple tensor and simulation backends
 License-Expression: MIT
 Project-URL: Homepage, https://github.com/UniEnvOrg/UniEnv

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

@@ -1,27 +1,27 @@
-unienv-0.0.1b4.dist-info/licenses/LICENSE,sha256=nkklvEaJUR4QDBygz7tkEe1FMVKV1JSjnGzJNLhdIWM,1091
+unienv-0.0.1b5.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/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/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=tdpKXzztd830FXzOE7SRYQ9Hu4cXmigy313EuF6Q-9c,12713
+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/dict_storage.py,sha256=DSqRIgo3m1XtUcLtyjYSqqpi01mr_nJOLg5BCddwPcg,13862
+unienv_data/storages/flattened.py,sha256=-1NBoCXSBMcE_kd5GRD_EvTyDupvpSzEdZka9DIWacU,7046
+unienv_data/storages/hdf5.py,sha256=uNhL7ji0Zzp4qRSXPS1H7Q74W1I4terUV8dNYKak06k,26730
+unienv_data/storages/pytorch.py,sha256=YczjJ9OVM-QqAYWuluSRSv80vZnsg1nP-b6vSk-6D2Y,6953
+unienv_data/storages/transformation.py,sha256=-9_jPZNpx6RXY_ojv_1UCSTa4Z9apI9V9jit8nG93oM,8133
+unienv_data/transformations/image_compress.py,sha256=pp9Q5CoVcZwu-oajnpFvtMuv7e3w2ZZV50S4pshoaj4,10000
 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,14 +36,14 @@ 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/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
@@ -54,20 +54,22 @@ unienv_interface/space/spaces/dict.py,sha256=G5_iYC1Bj5DqeJ7aFlq6eRJbnpATbIRIyRu
 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/filter_dict.py,sha256=DzR-hgHoHJObTipxwB2UrKVlTxbfIrJohaOgqdAICLY,5871
+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=fM5ukWUvNvPeDO48_PRU0KyyvGhBIDxaN9XZyQ1VaQQ,4364
 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
@@ -82,12 +84,12 @@ unienv_interface/wrapper/backend_compat.py,sha256=T6hosgu2hrZvg3xtnyELmR6Exlz-zt
 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.1b5.dist-info/METADATA,sha256=IH2-ZqP73SKp1QKGX05k16X0M5bl_JGWrSulUXdrkE4,3033
+unienv-0.0.1b5.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+unienv-0.0.1b5.dist-info/top_level.txt,sha256=wfcJ5_DruUtOEUZjEyfadaKn7B90hWqz2aw-eM3wX5g,29
+unienv-0.0.1b5.dist-info/RECORD,,

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/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,9 @@
 import abc
 import os
 import dataclasses
+import multiprocessing as mp
+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 +54,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 +62,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 +73,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 +81,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
@@ -97,6 +102,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 +124,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:
+            assert storage.is_multiprocessing_safe, "Storage is not multiprocessing safe"
+            self._lock = mp.Lock()
+            self._count_value = mp.Value("q", int(count))
+            self._offset_value = mp.Value("q", 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 self._multiprocessing else self._count_value
+
+    @count.setter
+    def count(self, value: int) -> None:
+        if self._multiprocessing:
+            self._count_value.value = int(value)
+        else:
+            self._count_value = int(value)
     
+    @property
+    def offset(self) -> int:
+        return self._offset_value.value if self._multiprocessing else self._offset_value
+
+    @offset.setter
+    def offset(self, value: int) -> None:
+        if self._multiprocessing:
+            self._offset_value.value = int(value)
+        else:
+            self._offset_value = int(value)
+
     @property
     def capacity(self) -> Optional[int]:
         return self.storage.capacity
@@ -176,12 +233,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 +261,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 +285,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 +302,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()