Replay Buffer Classes and Configuration in TorchRL

ReplayBuffer
├── Storage  — holds the raw data (list, tensor, memory-mapped, ...)
├── Sampler  — selects which indices to return (uniform, prioritized, slice, ...)
├── Writer   — decides where new data lands (round-robin, max-value, immutable, ...)
└── Transform — post-processes samples before returning them to the caller

from torchrl.data import (
    ReplayBuffer,
    TensorDictReplayBuffer,
    LazyTensorStorage,
    PrioritizedSampler,
    RoundRobinWriter,
)
from torchrl.envs import Compose, ObservationNorm

rb = TensorDictReplayBuffer(
    storage=LazyTensorStorage(max_size=100_000),
    sampler=PrioritizedSampler(max_capacity=100_000, alpha=0.6, beta=0.4),
    writer=RoundRobinWriter(),
    batch_size=256,
)

import torch
from tensordict import TensorDict
from torchrl.data import TensorDictPrioritizedReplayBuffer, LazyTensorStorage

rb = TensorDictPrioritizedReplayBuffer(
    alpha=0.7,
    beta=1.1,
    storage=LazyTensorStorage(max_size=10_000),
    batch_size=256,
)

data = TensorDict({"obs": torch.randn(1000, 8), "action": torch.randn(1000, 2)}, [1000])
rb.extend(data)

sample = rb.sample()
# sample["index"] contains buffer indices
# sample["priority_weight"] contains IS weights

# Update priorities after computing TD-errors:
sample.set("td_error", torch.rand(sample.shape))
rb.update_tensordict_priority(sample)

from torchrl.data import HindsightStrategy

# Available strategies:
HindsightStrategy.FUTURE   # Sample goal from a future step in the same episode (recommended)
HindsightStrategy.FINAL    # Use the final achieved state as goal
HindsightStrategy.EPISODE  # Sample any achieved state from the same episode uniformly
HindsightStrategy.RANDOM   # Sample a random achieved state from the entire buffer

import torch
from tensordict import TensorDict
from torchrl.data import HERReplayBuffer, HindsightStrategy, LazyTensorStorage

def my_reward_fn(td):
    # Compute sparse reward: 1 if achieved_goal ≈ desired_goal
    dist = (td["achieved_goal"] - td["desired_goal"]).norm(dim=-1, keepdim=True)
    return (dist < 0.05).float()

rb = HERReplayBuffer(
    reward_fn=my_reward_fn,
    her_ratio=0.8,
    strategy=HindsightStrategy.FUTURE,
    storage=LazyTensorStorage(max_size=50_000),
    batch_size=256,
)

from torchrl.data import LazyMemmapStorage, TensorDictReplayBuffer

# 10-million-step buffer backed by disk — no RAM limit
rb = TensorDictReplayBuffer(
    storage=LazyMemmapStorage(
        max_size=10_000_000,
        scratch_dir="/scratch/my_experiment/rb",
    ),
    batch_size=512,
)

from torchrl.data import RandomSampler, SamplerWithoutReplacement

# These two are equivalent:
s1 = RandomSampler(replacement=False)
s2 = SamplerWithoutReplacement()
isinstance(s1, SamplerWithoutReplacement)  # True

import torch
from tensordict import TensorDict
from torchrl.data import (
    TensorDictPrioritizedReplayBuffer,
    LazyTensorStorage,
)
from torchrl.envs import GymEnv, TransformedEnv, StepCounter

# --- Environment ---
env = TransformedEnv(GymEnv("CartPole-v1"), StepCounter())

# --- Replay buffer ---
rb = TensorDictPrioritizedReplayBuffer(
    alpha=0.7,
    beta=0.5,
    storage=LazyTensorStorage(max_size=50_000),
    batch_size=256,
)

# --- Warm-up: fill buffer with random transitions ---
td = env.reset()
for _ in range(1_000):
    td = env.rand_step(td)
    rb.add(td.clone())
    if td["next", "done"].any():
        td = env.reset()

# --- Training loop ---
for step in range(10_000):
    # Collect one transition
    td = env.rand_step(td)
    rb.add(td.clone())
    if td["next", "done"].any():
        td = env.reset()

    # Sample and train
    sample = rb.sample()
    # ... compute loss, td_errors ...
    td_errors = torch.rand(sample.shape)  # placeholder
    sample.set("td_error", td_errors)
    rb.update_tensordict_priority(sample)

    # Anneal beta toward 1.0
    rb._sampler._beta = min(1.0, 0.5 + step / 10_000 * 0.5)

from pathlib import Path

# Save the full buffer state (storage + sampler + writer metadata)
rb.dumps(Path("./checkpoints/replay_buffer"))

# Restore from disk
rb.loads(Path("./checkpoints/replay_buffer"))