[BugFix] Fix property getter in RayReplayBuffer (#2869)

Author: Vincent Moens

torchrl/data/replay_buffers/ray_buffer.py

@@ -207,3 +207,11 @@ def state_dict(self):
 
     def __len__(self):
         return ray.get(self._rb.__len__.remote())
+
+    @property
+    def write_count(self):
+        return ray.get(self._rb._getattr.remote("write_count"))
+
+    @property
+    def dim_extend(self):
+        return ray.get(self._rb._getattr.remote("dim_extend"))

torchrl/data/replay_buffers/replay_buffers.py

@@ -447,6 +447,10 @@ def __len__(self) -> int:
         with self._replay_lock:
             return len(self._storage)
 
+    def _getattr(self, attr):
+        # To access properties in remote settings, see RayReplayBuffer.write_count for instance
+        return getattr(self, attr)
+
     @property
     def write_count(self):
         """The total number of items written so far in the buffer through add and extend."""