From 66c3672738d1a2280af887cb03e04a6af8333ed4 Mon Sep 17 00:00:00 2001 From: s1lent4gnt Date: Mon, 31 Mar 2025 09:43:40 +0200 Subject: [PATCH] Fix: Prevent Invalid next_state References When optimize_memory=True (#918) --- lerobot/scripts/server/buffer.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/lerobot/scripts/server/buffer.py b/lerobot/scripts/server/buffer.py index 20787568..776ad9ec 100644 --- a/lerobot/scripts/server/buffer.py +++ b/lerobot/scripts/server/buffer.py @@ -286,9 +286,10 @@ class ReplayBuffer: raise RuntimeError("Cannot sample from an empty buffer. Add transitions first.") batch_size = min(batch_size, self.size) + high = max(0, self.size - 1) if self.optimize_memory and self.size < self.capacity else self.size # Random indices for sampling - create on the same device as storage - idx = torch.randint(low=0, high=self.size, size=(batch_size,), device=self.storage_device) + idx = torch.randint(low=0, high=high, size=(batch_size,), device=self.storage_device) # Identify image keys that need augmentation image_keys = [k for k in self.states if k.startswith("observation.image")] if self.use_drq else []