Offline RL Dataset Loaders and Experience Replay Classes

from torchrl.data.datasets import AtariDQNExperienceReplay

# Load run 1 of Pong (downloads ~10 GB on first use)
data = AtariDQNExperienceReplay(
    dataset_id="Pong/1",
    batch_size=32,
    num_slices=8,   # 8 non-overlapping sub-sequences per batch
)

sample = data.sample()
print(sample["observation"].shape)  # torch.Size([32, 1, 84, 84]) — grayscale frames

from torchrl.data.datasets import MinariExperienceReplay

data = MinariExperienceReplay("door-human-v1", batch_size=32, download="force")

for batch in data:
    # batch keys: observation, action,
    #              next/observation, next/reward, next/done, ...
    print(batch["action"].shape)  # e.g. torch.Size([32, 28])
    break

from torchrl.data.datasets.lerobot import lerobot_columns_to_tensordict
import torch

columns = {
    "observation.state": torch.zeros(4, 7),
    "observation.images.top": torch.zeros(4, 3, 8, 8, dtype=torch.uint8),
    "action": torch.zeros(4, 7),
    "episode_index": torch.tensor([0, 0, 1, 1]),
    "task": ["pick", "pick", "place", "place"],
}

td = lerobot_columns_to_tensordict(columns)
print(td["observation", "state"].shape)     # torch.Size([4, 7])
print(td["observation", "image", "top"].shape)  # torch.Size([4, 3, 8, 8])
print(td.get("language_instruction").tolist())  # ['pick', 'pick', 'place', 'place']

from torchrl.data.datasets import OpenXExperienceReplay

data = OpenXExperienceReplay(
    "fractal20220817_data",
    batch_size=32,
    root="/data/openx",
)

sample = data.sample()
print(sample.keys())

from torchrl.data import (
    OfflineToOnlineReplayBuffer,
    TensorDictReplayBuffer,
    LazyTensorStorage,
)
from torchrl.data.datasets import D4RLExperienceReplay

offline = D4RLExperienceReplay("hopper-medium-v2", batch_size=128)
online  = TensorDictReplayBuffer(
    storage=LazyTensorStorage(max_size=100_000),
    batch_size=128,
)

# 50% offline data, 50% online data
mixed = OfflineToOnlineReplayBuffer(
    replay_buffer=offline,
    storage=online._storage,
    batch_size=256,
)

from torchrl.data.datasets import D4RLExperienceReplay
from torchrl.envs import ObservationNorm, RewardScaling, Compose

data = D4RLExperienceReplay("hopper-medium-v2", batch_size=256)
data.append_transform(
    Compose(
        ObservationNorm(loc=0.0, scale=1.0, in_keys=["observation"]),
        RewardScaling(loc=0.0, scale=0.01),
    )
)

sample = data.sample()
print(sample["observation"].mean())   # ≈ 0
print(sample["next", "reward"].abs().max())  # scaled