torchrl 0.11.0__cp314-cp314-macosx_11_0_arm64.whl

torchrl/data/__init__.py

@@ -0,0 +1,187 @@
+# 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 .llm import (
+    AdaptiveKLController,
+    ConstantKLController,
+    ContentBase,
+    create_infinite_iterator,
+    get_dataloader,
+    History,
+    PairwiseDataset,
+    PromptData,
+    PromptTensorDictTokenizer,
+    RewardData,
+    RolloutFromModel,
+    TensorDictTokenizer,
+    TokenizedDatasetLoader,
+    TopKRewardSelector,
+)
+from .map import (
+    BinaryToDecimal,
+    HashToInt,
+    MCTSForest,
+    QueryModule,
+    RandomProjectionHash,
+    SipHash,
+    TensorDictMap,
+    TensorMap,
+    Tree,
+)
+from .postprocs import DensifyReward, MultiStep
+from .replay_buffers import (
+    CompressedListStorage,
+    CompressedListStorageCheckpointer,
+    Flat2TED,
+    FlatStorageCheckpointer,
+    H5Combine,
+    H5Split,
+    H5StorageCheckpointer,
+    ImmutableDatasetWriter,
+    LazyMemmapStorage,
+    LazyStackStorage,
+    LazyTensorStorage,
+    ListStorage,
+    ListStorageCheckpointer,
+    Nested2TED,
+    NestedStorageCheckpointer,
+    PrioritizedReplayBuffer,
+    PrioritizedSampler,
+    PrioritizedSliceSampler,
+    RandomSampler,
+    RayReplayBuffer,
+    RemoteTensorDictReplayBuffer,
+    ReplayBuffer,
+    ReplayBufferEnsemble,
+    RoundRobinWriter,
+    SamplerEnsemble,
+    SamplerWithoutReplacement,
+    SliceSampler,
+    SliceSamplerWithoutReplacement,
+    Storage,
+    StorageCheckpointerBase,
+    StorageEnsemble,
+    StorageEnsembleCheckpointer,
+    TED2Flat,
+    TED2Nested,
+    TensorDictMaxValueWriter,
+    TensorDictPrioritizedReplayBuffer,
+    TensorDictReplayBuffer,
+    TensorDictRoundRobinWriter,
+    TensorStorage,
+    TensorStorageCheckpointer,
+    Writer,
+    WriterEnsemble,
+)
+from .tensor_specs import (
+    Binary,
+    Bounded,
+    BoundedContinuous,
+    Categorical,
+    Choice,
+    Composite,
+    DEVICE_TYPING,
+    MultiCategorical,
+    MultiOneHot,
+    NonTensor,
+    OneHot,
+    Stacked,
+    StackedComposite,
+    TensorSpec,
+    Unbounded,
+    UnboundedContinuous,
+    UnboundedDiscrete,
+)
+from .utils import check_no_exclusive_keys, consolidate_spec, contains_lazy_spec
+
+__all__ = [
+    "AdaptiveKLController",
+    "Binary",
+    "BinaryToDecimal",
+    "Bounded",
+    "BoundedContinuous",
+    "Categorical",
+    "Choice",
+    "Composite",
+    "CompressedListStorage",
+    "CompressedListStorageCheckpointer",
+    "ConstantKLController",
+    "ContentBase",
+    "DEVICE_TYPING",
+    "DensifyReward",
+    "Flat2TED",
+    "FlatStorageCheckpointer",
+    "H5Combine",
+    "H5Split",
+    "H5StorageCheckpointer",
+    "HashToInt",
+    "History",
+    "ImmutableDatasetWriter",
+    "LazyMemmapStorage",
+    "LazyStackStorage",
+    "LazyTensorStorage",
+    "ListStorage",
+    "ListStorageCheckpointer",
+    "MCTSForest",
+    "MultiCategorical",
+    "MultiOneHot",
+    "MultiStep",
+    "Nested2TED",
+    "NestedStorageCheckpointer",
+    "NonTensor",
+    "OneHot",
+    "PairwiseDataset",
+    "PrioritizedReplayBuffer",
+    "PrioritizedSampler",
+    "PrioritizedSliceSampler",
+    "PromptData",
+    "PromptTensorDictTokenizer",
+    "QueryModule",
+    "RandomProjectionHash",
+    "RandomSampler",
+    "RayReplayBuffer",
+    "RemoteTensorDictReplayBuffer",
+    "ReplayBuffer",
+    "ReplayBufferEnsemble",
+    "RewardData",
+    "RolloutFromModel",
+    "RoundRobinWriter",
+    "SamplerEnsemble",
+    "SamplerWithoutReplacement",
+    "SipHash",
+    "SliceSampler",
+    "SliceSamplerWithoutReplacement",
+    "Stacked",
+    "StackedComposite",
+    "Storage",
+    "StorageCheckpointerBase",
+    "StorageEnsemble",
+    "StorageEnsembleCheckpointer",
+    "TED2Flat",
+    "TED2Nested",
+    "TensorDictMap",
+    "TensorDictMaxValueWriter",
+    "TensorDictPrioritizedReplayBuffer",
+    "TensorDictReplayBuffer",
+    "TensorDictRoundRobinWriter",
+    "TensorDictTokenizer",
+    "TensorMap",
+    "TensorSpec",
+    "TensorStorage",
+    "TensorStorageCheckpointer",
+    "TokenizedDatasetLoader",
+    "TopKRewardSelector",
+    "Tree",
+    "Unbounded",
+    "UnboundedContinuous",
+    "UnboundedDiscrete",
+    "Writer",
+    "WriterEnsemble",
+    "check_no_exclusive_keys",
+    "consolidate_spec",
+    "contains_lazy_spec",
+    "create_infinite_iterator",
+    "get_dataloader",
+]

torchrl/data/datasets/__init__.py

@@ -0,0 +1,58 @@
+# 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
+
+# Import classes that don't have external dependencies
+from .atari_dqn import AtariDQNExperienceReplay
+from .common import BaseDatasetExperienceReplay
+
+# Conditional imports for classes with external dependencies
+try:
+    from .d4rl import D4RLExperienceReplay
+except ImportError:
+    pass
+
+try:
+    from .minari_data import MinariExperienceReplay
+except ImportError:
+    pass
+
+try:
+    from .gen_dgrl import GenDGRLExperienceReplay
+except ImportError:
+    pass
+
+try:
+    from .openml import OpenMLExperienceReplay
+except ImportError:
+    pass
+
+try:
+    from .openx import OpenXExperienceReplay
+except ImportError:
+    pass
+
+try:
+    from .roboset import RobosetExperienceReplay
+except ImportError:
+    pass
+
+try:
+    from .vd4rl import VD4RLExperienceReplay
+except ImportError:
+    pass
+
+__all__ = [
+    "AtariDQNExperienceReplay",
+    "BaseDatasetExperienceReplay",
+    "D4RLExperienceReplay",
+    "MinariExperienceReplay",
+    "GenDGRLExperienceReplay",
+    "OpenMLExperienceReplay",
+    "OpenXExperienceReplay",
+    "RobosetExperienceReplay",
+    "VD4RLExperienceReplay",
+]