Repository: pytorch/rl

Issue:
[BUG] Unexpected behavior of SumSegmentTree Resulting in Invalid Slices in PrioritizedSliceSampler.sample()
In `PrioritizedSampler.sample()`, `_sum_tree.scan_lower_bound()` sometimes generates index greater than len(storage).

The test code is as follows:

```python
import torch
from torchrl._torchrl import SumSegmentTreeFp32
import numpy as np

def test_sum_tree():
    torch.manual_seed(0)
    np.random.seed(0)
    sum_tree = SumSegmentTreeFp32(500)
    
    # repeat to ensure the bug happens
    for _ in range(1000):
        # update priority
        index = torch.arange(0,100, dtype=torch.long, device=torch.device("cpu"))
        priority = torch.rand(100, device=torch.device("cpu"))
        sum_tree[index] = priority+1e-8 # 1e-8 are not essential, w/o 1e-8, bug still happens
        
        # sample
        p_sum = sum_tree.query(0, 100)
        mass = np.random.uniform(0.0, p_sum, size=1000000) # sample a lot to ensure the bug happens.
        scanned_index = sum_tree.scan_lower_bound(mass)
        if scanned_index.max()>=100:
            print("Unexpected index! p_sum:{}, mass.max():{}, scanned_index.max():{}".format(p_sum, mass.max(),scanned_index.max()))
        
test_sum_tree()

```

In `PrioritizedSampler.sample()`, this unexpected behavior is handled by `index.clamp_max_(len(storage) - 1)`, but it still causes unexpected behavior in `PrioritizedSliceSampler.sample()`.

This occurs because, although the sample at `index = len(storage) - 1` is covered by `preceding_stop_idx`, it can still be sampled as start index, leading to cross-trajectory sampling.

## Checklist

- [x] I have checked that there is no similar issue in the repo (**required**)
- [x] I have read the [documentation](https://github.com/pytorch/rl/tree/main/docs/) (**required**)
- [x] I have provided a minimal working example to reproduce the bug (**required**)



Discussion:
> This occurs because, although the sample at index = len(storage) - 1 is covered by preceding_stop_idx, it can still be sampled as start index, leading to cross-trajectory sampling.

Is this also the case with the patch in #2228?
How can the start be equal to len(storage)-1 if this is masked?
It's a new issue due to `_sum_tree`, I think it can't be solved with #2228.

In `PrioritizedSampler.sample()` we have:
```python
def sample(...)
    ...
    index = self._sum_tree.scan_lower_bound(mass)
    ...
    index.clamp_max_(len(storage) - 1)
    weight = torch.as_tensor(self._sum_tree[index])
    ...
    return index, {"_weight": weight}
```

For some unknown bug in `_sum_tree`, we randomly get an index greater than len(storage), then it's clamped to len(storage) - 1, which is why we get the wrong start index.
cc @xiaomengy any idea why we get out-of-bound samples from the PRB?

Resolving patch:
diff --git a/sota-implementations/cql/cql_offline.py b/sota-implementations/cql/cql_offline.py
--- a/sota-implementations/cql/cql_offline.py
+++ b/sota-implementations/cql/cql_offline.py
@@ -162,7 +162,8 @@ def main(cfg: "DictConfig"):  # noqa: F821
 
     pbar.close()
     torchrl_logger.info(f"Training time: {time.time() - start_time}")
-    eval_env.close()
+    if not eval_env.is_closed:
+        eval_env.close()
 
 
 if __name__ == "__main__":
diff --git a/sota-implementations/cql/cql_online.py b/sota-implementations/cql/cql_online.py
--- a/sota-implementations/cql/cql_online.py
+++ b/sota-implementations/cql/cql_online.py
@@ -227,8 +227,10 @@ def main(cfg: "DictConfig"):  # noqa: F821
     torchrl_logger.info(f"Training took {execution_time:.2f} seconds to finish")
 
     collector.shutdown()
-    eval_env.close()
-    train_env.close()
+    if not eval_env.is_closed:
+        eval_env.close()
+    if not train_env.is_closed:
+        train_env.close()
 
 
 if __name__ == "__main__":
diff --git a/torchrl/data/replay_buffers/samplers.py b/torchrl/data/replay_buffers/samplers.py
--- a/torchrl/data/replay_buffers/samplers.py
+++ b/torchrl/data/replay_buffers/samplers.py
@@ -475,6 +475,15 @@ def sample(self, storage: Storage, batch_size: int) -> torch.Tensor:
             index = index.unsqueeze(0)
         index.clamp_max_(len(storage) - 1)
         weight = torch.as_tensor(self._sum_tree[index])
+        # get indices where weight is 0
+        zero_weight = weight == 0
+        index = index
+        while zero_weight.any():
+            index = torch.where(zero_weight, index - 1, index)
+            if (index < 0).any():
+                raise RuntimeError("Failed to find a suitable index")
+            zero_weight = torch.as_tensor(self._sum_tree[index])
+            zero_weight = weight == 0
 
         # Importance sampling weight formula:
         #   w_i = (p_i / sum(p) * N) ^ (-beta)


Test patch: