torchrl 0.11.0__cp314-cp314-macosx_11_0_arm64.whl

torchrl/trainers/algorithms/configs/collectors.py

@@ -0,0 +1,216 @@
+# 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 functools import partial
+from typing import Any
+
+from omegaconf import MISSING
+
+from torchrl.trainers.algorithms.configs.common import ConfigBase
+from torchrl.trainers.algorithms.configs.envs import EnvConfig
+
+
+@dataclass
+class BaseCollectorConfig(ConfigBase):
+    """Parent class to configure a data collector."""
+
+
+@dataclass
+class CollectorConfig(BaseCollectorConfig):
+    """A class to configure a synchronous data collector (Collector)."""
+
+    create_env_fn: ConfigBase = MISSING
+    policy: Any = None
+    policy_factory: Any = None
+    frames_per_batch: int | None = None
+    total_frames: int = -1
+    init_random_frames: int | None = 0
+    device: str | None = None
+    storing_device: str | None = None
+    policy_device: str | None = None
+    env_device: str | None = None
+    create_env_kwargs: dict | None = None
+    max_frames_per_traj: int | None = None
+    reset_at_each_iter: bool = False
+    postproc: Any = None
+    split_trajs: bool = False
+    exploration_type: str = "RANDOM"
+    return_same_td: bool = False
+    interruptor: Any = None
+    set_truncated: bool = False
+    use_buffers: bool = False
+    replay_buffer: Any = None
+    extend_buffer: bool = False
+    trust_policy: bool = True
+    compile_policy: Any = None
+    cudagraph_policy: Any = None
+    no_cuda_sync: bool = False
+    weight_updater: Any = None
+    weight_sync_schemes: Any = None
+    track_policy_version: bool = False
+    local_init_rb: bool = False
+    _target_: str = "torchrl.collectors.Collector"
+    _partial_: bool = False
+
+    def __post_init__(self):
+        self.create_env_fn._partial_ = True
+        if self.policy_factory is not None:
+            self.policy_factory._partial_ = True
+        if self.weight_updater is not None:
+            self.weight_updater._partial_ = True
+
+
+# Legacy alias
+SyncDataCollectorConfig = CollectorConfig
+
+
+@dataclass
+class AsyncCollectorConfig(BaseCollectorConfig):
+    """Configuration for asynchronous data collector (AsyncCollector)."""
+
+    create_env_fn: ConfigBase = field(
+        default_factory=partial(EnvConfig, _partial_=True)
+    )
+    policy: Any = None
+    policy_factory: Any = None
+    frames_per_batch: int | None = None
+    init_random_frames: int | None = 0
+    total_frames: int = -1
+    device: str | None = None
+    storing_device: str | None = None
+    policy_device: str | None = None
+    env_device: str | None = None
+    create_env_kwargs: dict | None = None
+    max_frames_per_traj: int | None = None
+    reset_at_each_iter: bool = False
+    postproc: ConfigBase | None = None
+    split_trajs: bool = False
+    exploration_type: str = "RANDOM"
+    set_truncated: bool = False
+    use_buffers: bool = False
+    replay_buffer: ConfigBase | None = None
+    extend_buffer: bool = False
+    trust_policy: bool = True
+    compile_policy: Any = None
+    cudagraph_policy: Any = None
+    no_cuda_sync: bool = False
+    weight_updater: Any = None
+    weight_sync_schemes: Any = None
+    track_policy_version: bool = False
+    local_init_rb: bool = False
+    _target_: str = "torchrl.collectors.AsyncCollector"
+    _partial_: bool = False
+
+    def __post_init__(self):
+        self.create_env_fn._partial_ = True
+        if self.policy_factory is not None:
+            self.policy_factory._partial_ = True
+        if self.weight_updater is not None:
+            self.weight_updater._partial_ = True
+
+
+# Legacy alias
+AsyncDataCollectorConfig = AsyncCollectorConfig
+
+
+@dataclass
+class MultiSyncCollectorConfig(BaseCollectorConfig):
+    """Configuration for multi-synchronous data collector (MultiSyncCollector)."""
+
+    create_env_fn: Any = MISSING
+    num_workers: int | None = None
+    policy: Any = None
+    policy_factory: Any = None
+    frames_per_batch: int | None = None
+    init_random_frames: int | None = 0
+    total_frames: int = -1
+    device: str | None = None
+    storing_device: str | None = None
+    policy_device: str | None = None
+    env_device: str | None = None
+    create_env_kwargs: dict | None = None
+    max_frames_per_traj: int | None = None
+    reset_at_each_iter: bool = False
+    postproc: ConfigBase | None = None
+    split_trajs: bool = False
+    exploration_type: str = "RANDOM"
+    set_truncated: bool = False
+    use_buffers: bool = False
+    replay_buffer: ConfigBase | None = None
+    extend_buffer: bool = False
+    trust_policy: bool = True
+    compile_policy: Any = None
+    cudagraph_policy: Any = None
+    no_cuda_sync: bool = False
+    weight_updater: Any = None
+    weight_sync_schemes: Any = None
+    track_policy_version: bool = False
+    local_init_rb: bool = False
+    _target_: str = "torchrl.collectors.MultiSyncCollector"
+    _partial_: bool = False
+
+    def __post_init__(self):
+        for env_cfg in self.create_env_fn:
+            env_cfg._partial_ = True
+        if self.policy_factory is not None:
+            self.policy_factory._partial_ = True
+        if self.weight_updater is not None:
+            self.weight_updater._partial_ = True
+
+
+# Legacy alias
+MultiSyncCollectorConfig = MultiSyncCollectorConfig
+
+
+@dataclass
+class MultiAsyncCollectorConfig(BaseCollectorConfig):
+    """Configuration for multi-asynchronous data collector (MultiAsyncCollector)."""
+
+    create_env_fn: Any = MISSING
+    num_workers: int | None = None
+    policy: Any = None
+    policy_factory: Any = None
+    frames_per_batch: int | None = None
+    init_random_frames: int | None = 0
+    total_frames: int = -1
+    device: str | None = None
+    storing_device: str | None = None
+    policy_device: str | None = None
+    env_device: str | None = None
+    create_env_kwargs: dict | None = None
+    max_frames_per_traj: int | None = None
+    reset_at_each_iter: bool = False
+    postproc: ConfigBase | None = None
+    split_trajs: bool = False
+    exploration_type: str = "RANDOM"
+    set_truncated: bool = False
+    use_buffers: bool = False
+    replay_buffer: ConfigBase | None = None
+    extend_buffer: bool = False
+    trust_policy: bool = True
+    compile_policy: Any = None
+    cudagraph_policy: Any = None
+    no_cuda_sync: bool = False
+    weight_updater: Any = None
+    weight_sync_schemes: Any = None
+    track_policy_version: bool = False
+    local_init_rb: bool = False
+    _target_: str = "torchrl.collectors.MultiAsyncCollector"
+    _partial_: bool = False
+
+    def __post_init__(self):
+        for env_cfg in self.create_env_fn:
+            env_cfg._partial_ = True
+        if self.policy_factory is not None:
+            self.policy_factory._partial_ = True
+        if self.weight_updater is not None:
+            self.weight_updater._partial_ = True
+
+
+# Legacy alias
+MultiAsyncCollectorConfig = MultiAsyncCollectorConfig

torchrl/trainers/algorithms/configs/common.py

@@ -0,0 +1,41 @@
+# 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 abc import ABC, abstractmethod
+from dataclasses import dataclass
+
+from omegaconf import DictConfig
+
+
+@dataclass
+class ConfigBase(ABC):
+    """Abstract base class for all configuration classes.
+
+    This class serves as the foundation for all configuration classes in the
+    configurable configuration system, providing a common interface and structure.
+    """
+
+    @abstractmethod
+    def __post_init__(self) -> None:
+        """Post-initialization hook for configuration classes."""
+
+
+@dataclass
+class Config:
+    """A flexible config that allows arbitrary fields."""
+
+    def __init__(self, **kwargs):
+        self._config = DictConfig(kwargs)
+
+    def __getattr__(self, name):
+        return getattr(self._config, name)
+
+    def __setattr__(self, name, value):
+        if name == "_config":
+            super().__setattr__(name, value)
+        else:
+            setattr(self._config, name, value)

torchrl/trainers/algorithms/configs/data.py

@@ -0,0 +1,308 @@
+# 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
+
+from omegaconf import MISSING
+
+from torchrl.trainers.algorithms.configs.common import ConfigBase
+
+
+@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."""
+
+
+@dataclass
+class RoundRobinWriterConfig(WriterConfig):
+    """Configuration for round-robin writer that distributes data across multiple storages."""
+
+    _target_: str = "torchrl.data.replay_buffers.RoundRobinWriter"
+    compilable: 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."""
+
+
+@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__()
+
+
+@dataclass
+class WriterEnsembleConfig(WriterConfig):
+    """Configuration for ensemble writer that combines multiple writers."""
+
+    _target_: str = "torchrl.data.replay_buffers.WriterEnsemble"
+    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):
+    """Configuration for TensorDict round-robin writer."""
+
+    _target_: str = "torchrl.data.replay_buffers.TensorDictRoundRobinWriter"
+    compilable: 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
+    traj_key: Any = None
+    ends: Any = None
+    trajectories: Any = None
+    cache_values: bool = False
+    truncated_key: Any = ("next", "truncated")
+    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"
+
+
+@dataclass
+class SliceSamplerWithoutReplacementConfig(SamplerConfig):
+    """Configuration for slice sampling without replacement."""
+
+    _target_: str = "torchrl.data.replay_buffers.SliceSamplerWithoutReplacement"
+    num_slices: int | None = None
+    slice_len: int | None = None
+    end_key: Any = None
+    traj_key: Any = None
+    ends: Any = None
+    trajectories: Any = None
+    cache_values: bool = False
+    truncated_key: Any = ("next", "truncated")
+    strict_length: bool = True
+    compile: Any = False
+    span: Any = False
+    use_gpu: Any = False
+
+
+@dataclass
+class SliceSamplerConfig(SamplerConfig):
+    """Configuration for slice sampling from replay buffer."""
+
+    _target_: str = "torchrl.data.replay_buffers.SliceSampler"
+    num_slices: int | None = None
+    slice_len: int | None = None
+    end_key: Any = None
+    traj_key: Any = None
+    ends: Any = None
+    trajectories: Any = None
+    cache_values: bool = False
+    truncated_key: Any = ("next", "truncated")
+    strict_length: bool = True
+    compile: Any = False
+    span: Any = False
+    use_gpu: Any = False
+
+
+@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"
+
+
+@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 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."""
+
+
+@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__()
+
+
+@dataclass
+class ListStorageConfig(StorageConfig):
+    """Configuration for list-based storage in replay buffer."""
+
+    _target_: str = "torchrl.data.replay_buffers.ListStorage"
+    max_size: int | None = None
+    compilable: bool = False
+
+
+@dataclass
+class StorageEnsembleWriterConfig(StorageConfig):
+    """Configuration for storage ensemble writer."""
+
+    _target_: str = "torchrl.data.replay_buffers.StorageEnsembleWriter"
+    writers: list[Any] = MISSING
+    transforms: list[Any] = MISSING
+
+
+@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
+
+
+@dataclass
+class StorageEnsembleConfig(StorageConfig):
+    """Configuration for storage ensemble."""
+
+    _target_: str = "torchrl.data.replay_buffers.StorageEnsemble"
+    storages: list[Any] = MISSING
+    transforms: list[Any] = MISSING
+
+
+@dataclass
+class LazyMemmapStorageConfig(StorageConfig):
+    """Configuration for lazy memory-mapped storage."""
+
+    _target_: str = "torchrl.data.replay_buffers.LazyMemmapStorage"
+    max_size: int | None = None
+    device: Any = None
+    ndim: int = 1
+    compilable: bool = False
+    shared_init: bool = False
+
+
+@dataclass
+class LazyTensorStorageConfig(StorageConfig):
+    """Configuration for lazy tensor storage."""
+
+    _target_: str = "torchrl.data.replay_buffers.LazyTensorStorage"
+    max_size: int | None = None
+    device: Any = None
+    ndim: int = 1
+    compilable: bool = False
+    shared_init: bool = False
+
+
+@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."""
+
+
+@dataclass
+class TensorDictReplayBufferConfig(ReplayBufferBaseConfig):
+    """Configuration for TensorDict-based replay buffer."""
+
+    _target_: str = "torchrl.data.replay_buffers.TensorDictReplayBuffer"
+    sampler: Any = None
+    storage: Any = None
+    writer: Any = None
+    transform: Any = None
+    batch_size: int | None = None
+
+    def __post_init__(self) -> None:
+        """Post-initialization hook for TensorDict replay buffer configurations."""
+        super().__post_init__()
+
+
+@dataclass
+class ReplayBufferConfig(ReplayBufferBaseConfig):
+    """Configuration for generic replay buffer."""
+
+    _target_: str = "torchrl.data.replay_buffers.ReplayBuffer"
+    sampler: Any = None
+    storage: Any = None
+    writer: Any = None
+    transform: Any = None
+    batch_size: int | None = None
+    shared: bool = False

torchrl/trainers/algorithms/configs/envs.py

@@ -0,0 +1,104 @@
+# 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
+
+from omegaconf import MISSING
+
+from torchrl.envs.common import EnvBase
+from torchrl.trainers.algorithms.configs.common import ConfigBase
+
+
+@dataclass
+class EnvConfig(ConfigBase):
+    """Base configuration class for environments."""
+
+    _partial_: bool = False
+
+    def __post_init__(self) -> None:
+        """Post-initialization hook for environment configurations."""
+        self._partial_ = False
+
+
+@dataclass
+class BatchedEnvConfig(EnvConfig):
+    """Configuration for batched environments."""
+
+    create_env_fn: Any = MISSING
+    num_workers: int = 1
+    create_env_kwargs: dict = field(default_factory=dict)
+    batched_env_type: str = "parallel"
+    device: str | None = None
+    # batched_env_type: Literal["parallel", "serial", "async"] = "parallel"
+    _target_: str = "torchrl.trainers.algorithms.configs.envs.make_batched_env"
+
+    def __post_init__(self) -> None:
+        """Post-initialization hook for batched environment configurations."""
+        super().__post_init__()
+        if hasattr(self.create_env_fn, "_partial_"):
+            self.create_env_fn._partial_ = True
+
+
+@dataclass
+class TransformedEnvConfig(EnvConfig):
+    """Configuration for transformed environments."""
+
+    base_env: Any = MISSING
+    transform: Any = None
+    cache_specs: bool = True
+    auto_unwrap: bool | None = None
+    _target_: str = "torchrl.envs.TransformedEnv"
+
+
+def make_batched_env(
+    create_env_fn, num_workers, batched_env_type="parallel", device=None, **kwargs
+):
+    """Create a batched environment.
+
+    Args:
+        create_env_fn: Function to create individual environments or environment instance.
+        num_workers: Number of worker environments.
+        batched_env_type: Type of batched environment (parallel, serial, async).
+        device: Device to place the batched environment on.
+        **kwargs: Additional keyword arguments.
+
+    Returns:
+        The created batched environment instance.
+    """
+    from torchrl.envs import AsyncEnvPool, ParallelEnv, SerialEnv
+
+    if create_env_fn is None:
+        raise ValueError("create_env_fn must be provided")
+
+    if num_workers is None:
+        raise ValueError("num_workers must be provided")
+
+    # If create_env_fn is a config object, create a lambda that instantiates it each time
+    if isinstance(create_env_fn, EnvBase):
+        # Already an instance (either instantiated config or actual env), wrap in lambda
+        env_instance = create_env_fn
+
+        def env_fn(env_instance=env_instance):
+            return env_instance
+
+    else:
+        env_fn = create_env_fn
+    assert callable(env_fn), env_fn
+
+    # Add device to kwargs if provided
+    if device is not None:
+        kwargs["device"] = device
+
+    if batched_env_type == "parallel":
+        return ParallelEnv(num_workers, env_fn, **kwargs)
+    elif batched_env_type == "serial":
+        return SerialEnv(num_workers, env_fn, **kwargs)
+    elif batched_env_type == "async":
+        return AsyncEnvPool([env_fn] * num_workers, **kwargs)
+    else:
+        raise ValueError(f"Unknown batched_env_type: {batched_env_type}")