Source code for torchrl.trainers.algorithms.configs.data
# Copyright (c) Meta Platforms, Inc. and affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
from __future__ import annotations
from dataclasses import dataclass, field
from typing import Any, Literal, TYPE_CHECKING
from omegaconf import MISSING
from torchrl.data.replay_buffers import WriterEnsemble
from torchrl.trainers.algorithms.configs.common import ConfigBase
if TYPE_CHECKING:
_ReplayServiceBackend = Literal["direct", "ray"]
_ReplayTransport = Literal["auto", "direct", "ray", "distributed"]
_SliceOutputLayout = Literal["flat", "batch_time"]
else:
# OmegaConf structured configs resolve these aliases at runtime and do not
# support Literal on all TorchRL-supported versions.
_ReplayServiceBackend = str
_ReplayTransport = str
_SliceOutputLayout = str
@dataclass
class WriterConfig(ConfigBase):
"""Base configuration class for replay buffer writers."""
_target_: str = "torchrl.data.replay_buffers.Writer"
def __post_init__(self) -> None:
"""Post-initialization hook for writer configurations."""
[docs]
@dataclass
class RoundRobinWriterConfig(WriterConfig):
"""Hydra configuration for :class:`~torchrl.data.RoundRobinWriter`."""
_target_: str = "torchrl.data.replay_buffers.RoundRobinWriter"
compilable: bool = False
track_generations: bool = False
def __post_init__(self) -> None:
"""Post-initialization hook for round-robin writer configurations."""
super().__post_init__()
@dataclass
class SamplerConfig(ConfigBase):
"""Base configuration class for replay buffer samplers."""
_target_: str = "torchrl.data.replay_buffers.Sampler"
def __post_init__(self) -> None:
"""Post-initialization hook for sampler configurations."""
[docs]
@dataclass
class RandomSamplerConfig(SamplerConfig):
"""Configuration for random sampling from replay buffer."""
_target_: str = "torchrl.data.replay_buffers.RandomSampler"
def __post_init__(self) -> None:
"""Post-initialization hook for random sampler configurations."""
super().__post_init__()
[docs]
@dataclass
class ConsumingSamplerConfig(SamplerConfig):
"""Hydra configuration for :class:`~torchrl.data.replay_buffers.ConsumingSampler`.
Every kwarg accepted by ``ConsumingSampler.__init__`` is exposed as a field
here.
"""
_target_: str = "torchrl.data.replay_buffers.ConsumingSampler"
max_sample_count: int = 1
def _make_writer_ensemble(writers: list[Any], p: Any = None) -> WriterEnsemble:
"""Build a writer ensemble from Hydra's keyword-based representation."""
del p # Kept for compatibility with the historical Config schema.
return WriterEnsemble(*writers)
@dataclass
class WriterEnsembleConfig(WriterConfig):
"""Configuration for ensemble writer that combines multiple writers."""
_target_: str = "torchrl.trainers.algorithms.configs.data._make_writer_ensemble"
writers: list[Any] = field(default_factory=list)
p: Any = None
@dataclass
class TensorDictMaxValueWriterConfig(WriterConfig):
"""Configuration for TensorDict max value writer."""
_target_: str = "torchrl.data.replay_buffers.TensorDictMaxValueWriter"
rank_key: Any = None
reduction: str = "sum"
@dataclass
class TensorDictRoundRobinWriterConfig(WriterConfig):
"""Hydra configuration for :class:`~torchrl.data.TensorDictRoundRobinWriter`."""
_target_: str = "torchrl.data.replay_buffers.TensorDictRoundRobinWriter"
compilable: bool = False
track_generations: bool = False
@dataclass
class ImmutableDatasetWriterConfig(WriterConfig):
"""Configuration for immutable dataset writer."""
_target_: str = "torchrl.data.replay_buffers.ImmutableDatasetWriter"
@dataclass
class SamplerEnsembleConfig(SamplerConfig):
"""Configuration for ensemble sampler that combines multiple samplers."""
_target_: str = "torchrl.data.replay_buffers.SamplerEnsemble"
samplers: list[Any] = field(default_factory=list)
p: Any = None
@dataclass
class PrioritizedSliceSamplerConfig(SamplerConfig):
"""Configuration for prioritized slice sampling from replay buffer."""
num_slices: int | None = None
slice_len: int | None = None
end_key: Any = None
end_keys: Any = None
traj_key: Any = None
ends: Any = None
trajectories: Any = None
cache_values: bool = False
truncated_key: Any = ("next", "truncated")
init_key: Any = "is_init"
strict_length: bool = True
compile: Any = False
span: Any = False
use_gpu: Any = False
max_capacity: int | None = None
alpha: float | None = None
beta: float | None = None
eps: float | None = None
reduction: str | None = None
_target_: str = "torchrl.data.replay_buffers.PrioritizedSliceSampler"
[docs]
@dataclass
class SliceSamplerWithoutReplacementConfig(SamplerConfig):
"""Hydra configuration for :class:`~torchrl.data.SliceSamplerWithoutReplacement`."""
_target_: str = "torchrl.data.replay_buffers.SliceSamplerWithoutReplacement"
num_slices: int | None = None
slice_len: int | None = None
end_key: Any = None
end_keys: Any = None
traj_key: Any = None
ends: Any = None
trajectories: Any = None
cache_values: bool = False
truncated_key: Any = ("next", "truncated")
init_key: Any = "is_init"
strict_length: bool = True
output_layout: _SliceOutputLayout = "flat"
slice_end_key: Any = ("collector", "slice_end")
time_dim_name: str | None = "time"
compile: Any = False
span: Any = False
use_gpu: Any = False
[docs]
@dataclass
class SliceSamplerConfig(SamplerConfig):
"""Hydra configuration for :class:`~torchrl.data.SliceSampler`."""
_target_: str = "torchrl.data.replay_buffers.SliceSampler"
num_slices: int | None = None
slice_len: int | None = None
end_key: Any = None
end_keys: Any = None
traj_key: Any = None
step_key: Any = "step_count"
fragmented: bool = False
ends: Any = None
trajectories: Any = None
cache_values: bool = False
truncated_key: Any = ("next", "truncated")
init_key: Any = "is_init"
strict_length: bool = True
pad_output: bool = False
output_layout: _SliceOutputLayout = "flat"
slice_end_key: Any = ("collector", "slice_end")
time_dim_name: str | None = "time"
compile: Any = False
span: Any = False
use_gpu: Any = False
[docs]
@dataclass
class StreamingSliceSamplerConfig(SamplerConfig):
"""Hydra configuration for :class:`~torchrl.data.StreamingSliceSampler`."""
_target_: str = "torchrl.data.replay_buffers.StreamingSliceSampler"
slice_len: int = MISSING
end_key: Any = None
end_keys: Any = None
traj_key: Any = None
cache_values: bool = False
truncated_key: Any = ("next", "truncated")
init_key: Any = "is_init"
strict_length: bool = True
pad_output: bool = False
compile: Any = False
span: Any = False
use_gpu: Any = False
[docs]
@dataclass
class PrioritizedSamplerConfig(SamplerConfig):
"""Configuration for prioritized sampling from replay buffer."""
max_capacity: int | None = None
alpha: float | None = None
beta: float | None = None
eps: float | None = None
reduction: str | None = None
_target_: str = "torchrl.data.replay_buffers.PrioritizedSampler"
[docs]
@dataclass
class SamplerWithoutReplacementConfig(SamplerConfig):
"""Configuration for sampling without replacement."""
_target_: str = "torchrl.data.replay_buffers.SamplerWithoutReplacement"
drop_last: bool = False
shuffle: bool = True
@dataclass
class SampleUnitConfig(ConfigBase):
"""Base configuration class for replay buffer sample units.
See also :class:`~torchrl.data.replay_buffers.SampleUnit`.
"""
_target_: str = "torchrl.data.replay_buffers.SampleUnit"
def __post_init__(self) -> None:
"""Post-initialization hook for sample unit configurations."""
@dataclass
class TransitionConfig(SampleUnitConfig):
"""Hydra configuration for :class:`~torchrl.data.replay_buffers.Transition`.
``Transition.__init__`` takes no arguments, so this config only carries
the instantiation target.
"""
_target_: str = "torchrl.data.replay_buffers.Transition"
@dataclass
class SequenceConfig(SampleUnitConfig):
"""Hydra configuration for :class:`~torchrl.data.replay_buffers.Sequence`.
Every kwarg accepted by ``Sequence.__init__`` is exposed as a field here
with the same default.
"""
_target_: str = "torchrl.data.replay_buffers.Sequence"
length: int = MISSING
episode_boundary: str = "pad"
done_key: Any = ("next", "done")
burn_in: int = 0
bootstrap: int = 0
dilation: int = 1
@dataclass
class StorageConfig(ConfigBase):
"""Base configuration class for replay buffer storage."""
_partial_: bool = False
_target_: str = "torchrl.data.replay_buffers.Storage"
def __post_init__(self) -> None:
"""Post-initialization hook for storage configurations."""
[docs]
@dataclass
class TensorStorageConfig(StorageConfig):
"""Configuration for tensor-based storage in replay buffer."""
_target_: str = "torchrl.data.replay_buffers.TensorStorage"
max_size: int | None = None
storage: Any = None
device: Any = None
ndim: int | None = None
compilable: bool = False
def __post_init__(self) -> None:
"""Post-initialization hook for tensor storage configurations."""
super().__post_init__()
[docs]
@dataclass
class ListStorageConfig(StorageConfig):
"""Hydra configuration for :class:`~torchrl.data.replay_buffers.ListStorage`.
Every kwarg accepted by ``ListStorage.__init__`` is exposed as a field here.
"""
_target_: str = "torchrl.data.replay_buffers.ListStorage"
max_size: int | None = None
compilable: bool = False
device: Any = None
[docs]
@dataclass
class StorageEnsembleWriterConfig(StorageConfig):
"""Hydra configuration for :class:`~torchrl.data.replay_buffers.WriterEnsemble`.
This name is a historical typo for
:class:`~torchrl.trainers.algorithms.configs.data.WriterEnsembleConfig`.
The Config is kept so existing Hydra group references do not vanish
without a deprecation cycle.
Fields match :class:`~torchrl.trainers.algorithms.configs.data.WriterEnsembleConfig`
(``writers``, ``p``). ``WriterEnsemble.__init__`` only accepts ``*writers``;
``p`` is stored for Config parity and is not a constructor argument.
"""
_target_: str = "torchrl.trainers.algorithms.configs.data._make_writer_ensemble"
writers: list[Any] = field(default_factory=list)
p: Any = None
[docs]
@dataclass
class LazyStackStorageConfig(StorageConfig):
"""Configuration for lazy stack storage."""
_target_: str = "torchrl.data.replay_buffers.LazyStackStorage"
max_size: int | None = None
compilable: bool = False
stack_dim: int = 0
[docs]
@dataclass
class StorageEnsembleConfig(StorageConfig):
"""Configuration for storage ensemble."""
_target_: str = "torchrl.data.replay_buffers.StorageEnsemble"
storages: list[Any] = MISSING
transforms: list[Any] = MISSING
[docs]
@dataclass
class LazyMemmapStorageConfig(StorageConfig):
"""Hydra configuration for :class:`~torchrl.data.replay_buffers.LazyMemmapStorage`.
Every kwarg accepted by ``LazyMemmapStorage.__init__`` is exposed as a field here.
"""
_target_: str = "torchrl.data.replay_buffers.LazyMemmapStorage"
max_size: int | None = None
scratch_dir: Any = None
device: Any = "cpu"
ndim: int = 1
existsok: bool = False
compilable: bool = False
shared_init: bool = False
auto_cleanup: bool | None = None
[docs]
@dataclass
class LazyTensorStorageConfig(StorageConfig):
"""Hydra configuration for :class:`~torchrl.data.replay_buffers.LazyTensorStorage`.
Every kwarg accepted by ``LazyTensorStorage.__init__`` is exposed as a field here.
"""
_target_: str = "torchrl.data.replay_buffers.LazyTensorStorage"
max_size: int | None = None
device: Any = "cpu"
ndim: int = 1
compilable: bool = False
consolidated: bool = False
shared_init: bool = False
cleanup_memmap: bool = True
@dataclass
class ReplayBufferBaseConfig(ConfigBase):
"""Base configuration class for replay buffers."""
_partial_: bool = False
def __post_init__(self) -> None:
"""Post-initialization hook for replay buffer configurations."""
[docs]
@dataclass
class TensorDictReplayBufferConfig(ReplayBufferBaseConfig):
"""Hydra configuration for :class:`~torchrl.data.replay_buffers.TensorDictReplayBuffer`.
Every kwarg accepted by ``TensorDictReplayBuffer.__init__`` (plus the ``ReplayBuffer``
kwargs it forwards via ``**kwargs``) is exposed as a field here.
"""
_target_: str = "torchrl.data.replay_buffers.TensorDictReplayBuffer"
priority_key: str = "td_error"
sampler: Any = None
sample_unit: Any = None
storage: Any = None
writer: Any = None
collate_fn: Any = None
pin_memory: bool = False
prefetch: int | None = None
transform: Any = None
transform_factory: Any = None
batch_size: int | None = None
dim_extend: int | None = None
checkpointer: Any = None
generator: Any = None
consume_after_n_samples: int | None = None
shared: bool = False
compilable: bool | None = None
delayed_init: bool | None = None
service_backend: _ReplayServiceBackend = "direct"
service_backend_options: dict[str, Any] = field(default_factory=dict)
def __post_init__(self) -> None:
"""Post-initialization hook for TensorDict replay buffer configurations."""
super().__post_init__()
[docs]
@dataclass
class ReplayBufferConfig(ReplayBufferBaseConfig):
"""Hydra configuration for :class:`~torchrl.data.replay_buffers.ReplayBuffer`.
Every kwarg accepted by ``ReplayBuffer.__init__`` is exposed as a field here.
"""
_target_: str = "torchrl.data.replay_buffers.ReplayBuffer"
storage: Any = None
sampler: Any = None
sample_unit: Any = None
writer: Any = None
collate_fn: Any = None
pin_memory: bool = False
prefetch: int | None = None
transform: Any = None
transform_factory: Any = None
batch_size: int | None = None
dim_extend: int | None = None
checkpointer: Any = None
generator: Any = None
consume_after_n_samples: int | None = None
shared: bool = False
compilable: bool | None = None
delayed_init: bool | None = None
service_backend: _ReplayServiceBackend = "direct"
service_backend_options: dict[str, Any] = field(default_factory=dict)
transport: _ReplayTransport = "auto"
transport_options: dict[str, Any] | None = None
