[BUG] SliceSampler doesn't work as expected when collecting data from parallel environment?

[briannnyee]

Hi TorchRL devs,

I currently have this settings (simplified for clarity):

sampler = SliceSampler(
            slice_len=4, 
            end_key=None,
            traj_key=("collector", "traj_ids"),
            truncated_key=None,
		    strict_length=True,
)

...some codes...

frames_per_batch = num_envs * num_steps_per_env
collector = SyncDataCollectorWrapper(
            create_env_fn=env,
            policy=actor_module,
            frames_per_batch=frames_per_batch,
            total_frames=total_frames,
            init_random_frames=init_random_frames,
            exploration_type=ExplorationType.RANDOM,
            device=self.device,
)

...some codes...

data = next(collector_iter)
self.replay_buffer.extend(data.reshape(-1))

batch = self.replay_buffer.sample()
# RuntimeError: Did not find a single trajectory with sufficient length (length range: 1 - 1 / required=4))

After spending some time investigating this, I realized that the problem could be because SliceSampler expects wrong format of traj_key. Let's say we have num_envs=2 and num_steps_per_env=1, SliceSampler expects the data is stored in a episodic way, e.g. traj_key=[0,0,0,...0,1,1,1...,1]. While in reality, the data is stored sequentially, traj_key=[0,1,0,1,...,0,1].

Did I do something wrong here? or is there a way to workaround this / is it a bug that needs a patch?

My torchrl version is 0.8. Let me know if I need to provide more info. Thanks in advance!

[vmoens]

Hello!
Have you tried passing the ndim=2 arg to your replay buffer storage?

There is a way to circumvent this by telling the storage how many dimensions it should take into account when saving data. This can be done through the ndim keyword argument which is accepted by all contiguous storages such as TensorStorage and the likes.

Then you won't have to call reshape(-1) and you will store each traj independently as it should be.

[briannnyee]

Hello!

I tried with ndim=2, num_envs=16, num_steps_per_env=1 and I got this error

self.replay_buffer.extend(data)
 # File "~/python/lib/python3.10/site-packages/torchrl/data/replay_buffers/replay_buffers.py", line 1364, in _set_index_in_td
 # tensordict.set("index", expand_as_right(index, tensordict))
 # File "~/python/lib/python3.10/site-packages/tensordict/utils.py", line 383, in expand_as_right
 #  raise RuntimeError(
 # RuntimeError: tensor shape is incompatible with dest shape, got: tensor.shape=torch.Size([16, 2]), dest=torch.Size([16, 1])

[vmoens]

Can you comment on how your use case differs from this?

from torchrl.envs import ParallelEnv, GymEnv, StepCounter
from torchrl.data import ReplayBuffer, LazyTensorStorage, SliceSampler
from functools import partial
import torch

def main():
    torch.manual_seed(0)
    make_env = partial(GymEnv, "Pendulum-v1")
    env_parallel = ParallelEnv(4, make_env).append_transform(StepCounter())
    env_parallel.set_seed(0)
    rb = ReplayBuffer(storage=LazyTensorStorage(1000, ndim=2), sampler=SliceSampler(num_slices=3))
    rb.extend(env_parallel.rollout(100, break_when_any_done=False))
    s = rb.sample(30)
    print(s["step_count"].squeeze())

if __name__ == "__main__":
    main()

That prints

tensor([87, 88, 89, 90, 91, 92, 93, 94, 95, 96, 64, 65, 66, 67, 68, 69, 70, 71,
        72, 73, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19])

[briannnyee]

I don't have StepCounter() transform. Also my env is IsaacLab.

[vmoens]

The step counter is just here to show that the sampler does its job.

Can you give me another copy of the (pseudo)code you're using to init the env / collector / replay buffer?

Also what is the shape of the tensordict you're using in extend the first and second time? Maybe you could copy what print(tensordict) gives you?

[briannnyee]

Can you give me another copy of the (pseudo)code you're using to init the env / collector / replay buffer?

I'm happy to and one point to add on my prev comment is I didn't use rollout():

(pseudocode)
env = IsaacLabWrapper(env)

sampler = SliceSampler(
                slice_len=4, 
                end_key=None,
                traj_key=("collector", "traj_ids"),
                truncated_key=None,
                strict_length=True,
)

replay_buffer_size = 1_000_000
buffer_batch_size = 256
rb = TensorDictReplayBuffer(
            storage=LazyTensorStorage(
                max_size=replay_buffer_size, device=device, ndim=2 # added this
            ),
            sampler=sampler,
            batch_size=buffer_batch_size,
            priority_key="td_error"
)

num_envs = 16
num_steps_per_env = 1
learning_starts = 10
frames_per_batch = num_envs * num_steps_per_env
init_random_frames = learning_starts * frames_per_batch
collector = SyncDataCollectorWrapper(
            create_env_fn=env,
            policy=actor_module,
            frames_per_batch=frames_per_batch,
            total_frames=total_frames,
            init_random_frames=init_random_frames,
            exploration_type=ExplorationType.RANDOM,
            device=device,
)

for data in collector:
    rb.extend(data)
    # RuntimeError: tensor shape is incompatible with dest shape, got: tensor.shape=torch.Size([16, 2]), dest=torch.Size([16, 1])

Also what is the shape of the tensordict you're using in extend the first and second time? Maybe you could copy what print(tensordict) gives you?

The first time was (num_envs * num_steps_per_env,)=(16,):

> data_reshaped = data.reshape(-1)
> print(data_reshaped)
TensorDict(
    fields={
        ...
        collector: TensorDict(
            fields={
                traj_ids: Tensor(shape=torch.Size([16]), device=cuda:0, dtype=torch.int64, is_shared=True)},
            batch_size=torch.Size([16]),
            device=cuda:0,
            is_shared=True),
        ...
    batch_size=torch.Size([16]),
    device=cuda:0,
    is_shared=True)
> print(data_reshaped["collector"]["traj_ids"])
tensor([ 0,  1,  2,  3,  4,  5,  6,  7,  8,  9, 10, 11, 12, 13, 14, 15],
       device='cuda:0')
> rb.extend(data_reshaped)
# No Error
> batch = rb.sample()
# RuntimeError: Did not find a single trajectory with sufficient length (length range: 1 - 1 / required=4))

and the second time was (num_envs, num_steps_per_env)=(16, 1):

> print(data)
TensorDict(
    fields={
       ...
        collector: TensorDict(
            fields={
                traj_ids: Tensor(shape=torch.Size([16, 1]), device=cuda:0, dtype=torch.int64, is_shared=True)},
            batch_size=torch.Size([16, 1]),
            device=cuda:0,
            is_shared=True),
       ...
    batch_size=torch.Size([16, 1]),
    device=cuda:0,
    is_shared=True)

> print(data["collector"]["traj_ids"])
tensor([[ 0],
        [ 1],
        [ 2],
        [ 3],
        [ 4],
        [ 5],
        [ 6],
        [ 7],
        [ 8],
        [ 9],
        [10],
        [11],
        [12],
        [13],
        [14],
        [15]], device='cuda:0')
> rb.extend(data)
# RuntimeError: tensor shape is incompatible with dest shape, got: tensor.shape=torch.Size([16, 2]), dest=torch.Size([16, 1])

[vmoens]

Here's a test with IsaacLab that works fine on my end

rl/test/test_libs.py

Line 5153 in 071d079

def test_isaaclab_rb(self, env):

Can you check what differs in your example?

I don't get how the buffer is initialized, what the collector wrapper does etc...