[BUG] Replay buffers on LazyMemmapStorage silently lose string and other non-tensor fields #4565

Describe the bug

A replay buffer on LazyMemmapStorage loses the strings and other non-tensor fields written to it, with no error or warning. Every row reads back the value of the first element of the first batch. Tensor fields are stored correctly, and LazyTensorStorage stores the same data correctly. Anyone who keeps text (prompts, instructions, mission strings) or other Python objects in a memory-mapped buffer trains on wrong data without knowing it.

To Reproduce

import tempfile

import torch
from tensordict import NonTensorStack, TensorDict
from torchrl.data import LazyMemmapStorage, LazyTensorStorage, ReplayBuffer


def batch(start):
    return TensorDict(
        obs=torch.arange(start, start + 3).float(),
        text=NonTensorStack(*[f"step {i}" for i in range(start, start + 3)]),
        batch_size=[3],
    )


for storage in (LazyTensorStorage(10), LazyMemmapStorage(10, scratch_dir=tempfile.mkdtemp())):
    rb = ReplayBuffer(storage=storage)
    rb.extend(batch(0))
    rb.extend(batch(3))
    out = rb[:6]
    print(type(storage).__name__, out["obs"].tolist(), list(out["text"]))

LazyTensorStorage [0.0, 1.0, 2.0, 3.0, 4.0, 5.0] ['step 0', 'step 1', 'step 2', 'step 3', 'step 4', 'step 5']
LazyMemmapStorage [0.0, 1.0, 2.0, 3.0, 4.0, 5.0] ['step 0', 'step 0', 'step 0', 'step 0', 'step 0', 'step 0']

TensorDictReplayBuffer gives the same result.

Expected behavior

LazyMemmapStorage returns step 0 to step 5, as LazyTensorStorage does.

System info

Installed with pip: torchrl 0.14.1, tensordict 0.14.3, torch 2.14.1+cpu, numpy 2.5.4, Python 3.12.12, linux. TorchRL main (f0af32d) with tensordict main gives the same output.

Reason and Possible fixes

A memory-mapped storage is locked once it is initialized. For a locked storage, TensorStorage.set keeps only the keys that _storage_keys lists (tensor.py#L277-L297, #L676-L683). _storage_keys lists them with keys(include_nested=True, leaves_only=True, is_leaf=_NESTED_TENSORS_AS_LISTS). That predicate does not count non-tensor entries as leaves, so text is not in the list, and data.select(*storage_keys, strict=False) drops it from every write. The storage keeps what it was initialized with: the first element, expanded to every row.

The filter came with #3323, to drop keys that the storage does not have. text is in the storage, so dropping it was not the intent.

Listing the keys with the public tensordict.is_leaf_nontensor predicate instead (is_leaf=is_leaf_nontensor, available in tensordict 0.14.3 and main) fixes the repro: with that one change, both storages return step 0 to step 5, on tensordict 0.14.3 and on main. _raise_informative_lock_error (#L813-L821) lists keys with the same predicate, so its message leaves out non-tensor keys too.

Checklist

• I have checked that there is no similar issue in the repo (required)
• I have read the documentation (required)
• I have provided a minimal working example to reproduce the bug (required)