[Feature] Offline datasets: D4RL

[vmoens]

Description

Integrates offline RL dataset in torchrl

>>> from torchrl.data.datasets import D4RLExperienceReplay
>>> data = D4RLExperienceReplay('kitchen-complete-v0', split_trajs=True)
>>> print(data._storage._storage)
TensorDict(
    fields={
        _batch_size: MemmapTensor(shape=torch.Size([20, 1]), device=cpu, dtype=torch.int64, is_shared=False),
        action: MemmapTensor(shape=torch.Size([20, 207, 9]), device=cpu, dtype=torch.float32, is_shared=False),
        done: MemmapTensor(shape=torch.Size([20, 207]), device=cpu, dtype=torch.bool, is_shared=False),
        index: MemmapTensor(shape=torch.Size([20]), device=cpu, dtype=torch.int32, is_shared=False),
        infos: MemmapTensor(shape=torch.Size([20, 207]), device=cpu, dtype=torch.int64, is_shared=False),
        mask: MemmapTensor(shape=torch.Size([20, 207]), device=cpu, dtype=torch.bool, is_shared=False),
        next: TensorDict(
            fields={
                done: MemmapTensor(shape=torch.Size([20, 207]), device=cpu, dtype=torch.bool, is_shared=False),
                observation: MemmapTensor(shape=torch.Size([20, 207, 60]), device=cpu, dtype=torch.float32, is_shared=False),
                reward: MemmapTensor(shape=torch.Size([20, 207]), device=cpu, dtype=torch.float32, is_shared=False)},
            batch_size=torch.Size([20, 207]),
            device=cpu,
            is_shared=False),
        observation: MemmapTensor(shape=torch.Size([20, 207, 60]), device=cpu, dtype=torch.float32, is_shared=False),
        timeouts: MemmapTensor(shape=torch.Size([20, 207]), device=cpu, dtype=torch.bool, is_shared=False),
        traj_ids: MemmapTensor(shape=torch.Size([20, 207]), device=cpu, dtype=torch.int64, is_shared=False)},
    batch_size=torch.Size([20]),
    device=cpu,
    is_shared=False)
>>> print(data.sample(10))  # will sample 10 trajectories since split_trajs is set to True
TensorDict(
    fields={
        action: Tensor(shape=torch.Size([10, 207, 9]), device=cpu, dtype=torch.float32, is_shared=False),
        done: Tensor(shape=torch.Size([10, 207]), device=cpu, dtype=torch.bool, is_shared=False),
        index: Tensor(shape=torch.Size([10, 207]), device=cpu, dtype=torch.int32, is_shared=False),
        infos: Tensor(shape=torch.Size([10, 207]), device=cpu, dtype=torch.int64, is_shared=False),
        mask: Tensor(shape=torch.Size([10, 207]), device=cpu, dtype=torch.bool, is_shared=False),
        next: TensorDict(
            fields={
                done: Tensor(shape=torch.Size([10, 207]), device=cpu, dtype=torch.bool, is_shared=False),
                observation: Tensor(shape=torch.Size([10, 207, 60]), device=cpu, dtype=torch.float32, is_shared=False),
                reward: Tensor(shape=torch.Size([10, 207]), device=cpu, dtype=torch.float32, is_shared=False)},
            batch_size=torch.Size([10, 207]),
            device=cpu,
            is_shared=False),
        observation: Tensor(shape=torch.Size([10, 207, 60]), device=cpu, dtype=torch.float32, is_shared=False),
        timeouts: Tensor(shape=torch.Size([10, 207]), device=cpu, dtype=torch.bool, is_shared=False),
        traj_ids: Tensor(shape=torch.Size([10, 207]), device=cpu, dtype=torch.int64, is_shared=False)},
    batch_size=torch.Size([10, 207]),
    device=cpu,
    is_shared=False)

These datasets can be used with transforms:

>>> from torchrl.data.datasets.d4rl import D4RLExperienceReplay
>>> from torchrl.envs import ObservationNorm
>>> data = D4RLExperienceReplay("maze2d-umaze-v1")
>>> # we can append transforms to the dataset
>>> data.append_transform(ObservationNorm(loc=-1, scale=1.0))
>>> data.sample(128)
TensorDict(
    fields={
        action: Tensor(shape=torch.Size([128, 2]), device=cpu, dtype=torch.float32, is_shared=False),
        done: Tensor(shape=torch.Size([128]), device=cpu, dtype=torch.bool, is_shared=False),
        index: Tensor(shape=torch.Size([128]), device=cpu, dtype=torch.int32, is_shared=False),
        next: TensorDict(
            fields={
                done: Tensor(shape=torch.Size([128]), device=cpu, dtype=torch.bool, is_shared=False),
                observation: Tensor(shape=torch.Size([128, 4]), device=cpu, dtype=torch.float32, is_shared=False)},
                reward: Tensor(shape=torch.Size([128]), device=cpu, dtype=torch.float32, is_shared=False)},
            batch_size=torch.Size([128]),
            device=cpu,
            is_shared=False),
        observation: Tensor(shape=torch.Size([128, 4]), device=cpu, dtype=torch.float32, is_shared=False),
    batch_size=torch.Size([128]),
    device=cpu,
    is_shared=False)