torchrl 0.11.0__cp314-cp314t-macosx_11_0_arm64.whl

torchrl/envs/vec_envs.py

@@ -0,0 +1,11 @@
+# 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 warnings
+
+warnings.warn("vec_env.py has moved to batch_envs.py.", category=DeprecationWarning)
+
+from .batched_envs import *  # noqa: F403, F401

torchrl/modules/__init__.py

@@ -0,0 +1,206 @@
+# 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 torchrl.modules.tensordict_module.common import DistributionalDQNnet
+
+from .distributions import (
+    Delta,
+    distributions_maps,
+    IndependentNormal,
+    LLMMaskedCategorical,
+    MaskedCategorical,
+    MaskedOneHotCategorical,
+    NormalParamExtractor,
+    NormalParamWrapper,
+    OneHotCategorical,
+    OneHotOrdinal,
+    Ordinal,
+    ReparamGradientStrategy,
+    TanhDelta,
+    TanhNormal,
+    TruncatedNormal,
+)
+from .models import (
+    BatchRenorm1d,
+    ConsistentDropout,
+    ConsistentDropoutModule,
+    Conv3dNet,
+    ConvNet,
+    DdpgCnnActor,
+    DdpgCnnQNet,
+    DdpgMlpActor,
+    DdpgMlpQNet,
+    DecisionTransformer,
+    DreamerActor,
+    DTActor,
+    DuelingCnnDQNet,
+    MLP,
+    MultiAgentConvNet,
+    MultiAgentMLP,
+    MultiAgentNetBase,
+    NoisyLazyLinear,
+    NoisyLinear,
+    ObsDecoder,
+    ObsEncoder,
+    OnlineDTActor,
+    QMixer,
+    reset_noise,
+    RSSMPosterior,
+    RSSMPrior,
+    RSSMRollout,
+    Squeeze2dLayer,
+    SqueezeLayer,
+    VDNMixer,
+)
+from .tensordict_module import (
+    Actor,
+    ActorCriticOperator,
+    ActorCriticWrapper,
+    ActorValueOperator,
+    AdditiveGaussianModule,
+    AdditiveGaussianWrapper,
+    DecisionTransformerInferenceWrapper,
+    DistributionalQValueActor,
+    DistributionalQValueHook,
+    DistributionalQValueModule,
+    EGreedyModule,
+    EGreedyWrapper,
+    GRU,
+    GRUCell,
+    GRUModule,
+    LMHeadActorValueOperator,
+    LSTM,
+    LSTMCell,
+    LSTMModule,
+    MultiStepActorWrapper,
+    OrnsteinUhlenbeckProcessModule,
+    OrnsteinUhlenbeckProcessWrapper,
+    ProbabilisticActor,
+    QValueActor,
+    QValueHook,
+    QValueModule,
+    recurrent_mode,
+    SafeModule,
+    SafeProbabilisticModule,
+    SafeProbabilisticTensorDictSequential,
+    SafeSequential,
+    set_recurrent_mode,
+    TanhModule,
+    ValueOperator,
+    VmapModule,
+    WorldModelWrapper,
+)
+from .tensordict_module.exploration import (
+    RandomPolicy,
+    set_exploration_modules_spec_from_env,
+)
+from .utils import get_primers_from_module
+from .planners import CEMPlanner, MPCPlannerBase, MPPIPlanner  # usort:skip
+from .mcts import (  # usort:skip
+    EXP3Score,
+    MCTSScore,
+    MCTSScores,
+    PUCTScore,
+    UCB1TunedScore,
+    UCBScore,
+)
+
+__all__ = [
+    "Actor",
+    "ActorCriticOperator",
+    "ActorCriticWrapper",
+    "ActorValueOperator",
+    "AdditiveGaussianModule",
+    "AdditiveGaussianWrapper",
+    "BatchRenorm1d",
+    "CEMPlanner",
+    "EXP3Score",
+    "ConsistentDropout",
+    "ConsistentDropoutModule",
+    "Conv3dNet",
+    "ConvNet",
+    "DTActor",
+    "DdpgCnnActor",
+    "DdpgCnnQNet",
+    "DdpgMlpActor",
+    "DdpgMlpQNet",
+    "DecisionTransformer",
+    "DecisionTransformerInferenceWrapper",
+    "Delta",
+    "DistributionalDQNnet",
+    "DistributionalQValueActor",
+    "DistributionalQValueHook",
+    "DistributionalQValueModule",
+    "DreamerActor",
+    "DuelingCnnDQNet",
+    "EGreedyModule",
+    "EGreedyWrapper",
+    "GRU",
+    "GRUCell",
+    "GRUModule",
+    "LLMMaskedCategorical",
+    "IndependentNormal",
+    "LMHeadActorValueOperator",
+    "LSTM",
+    "LSTMCell",
+    "LSTMModule",
+    "MCTSScore",
+    "MCTSScores",
+    "MLP",
+    "MPCPlannerBase",
+    "MPPIPlanner",
+    "MaskedCategorical",
+    "MaskedOneHotCategorical",
+    "MultiAgentConvNet",
+    "MultiAgentMLP",
+    "LLMMaskedCategorical",
+    "MultiAgentNetBase",
+    "MultiStepActorWrapper",
+    "NoisyLazyLinear",
+    "NoisyLinear",
+    "NormalParamExtractor",
+    "NormalParamWrapper",
+    "ObsDecoder",
+    "ObsEncoder",
+    "OneHotCategorical",
+    "OneHotOrdinal",
+    "OnlineDTActor",
+    "Ordinal",
+    "OrnsteinUhlenbeckProcessModule",
+    "OrnsteinUhlenbeckProcessWrapper",
+    "ProbabilisticActor",
+    "PUCTScore",
+    "QMixer",
+    "QValueActor",
+    "QValueHook",
+    "QValueModule",
+    "RSSMPosterior",
+    "RSSMPrior",
+    "RSSMRollout",
+    "ReparamGradientStrategy",
+    "SafeModule",
+    "SafeProbabilisticModule",
+    "SafeProbabilisticTensorDictSequential",
+    "SafeSequential",
+    "Squeeze2dLayer",
+    "SqueezeLayer",
+    "TanhDelta",
+    "TanhModule",
+    "TanhNormal",
+    "TruncatedNormal",
+    "UCB1TunedScore",
+    "UCBScore",
+    "VDNMixer",
+    "ValueOperator",
+    "VmapModule",
+    "WorldModelWrapper",
+    "distributions_maps",
+    "get_primers_from_module",
+    "recurrent_mode",
+    "reset_noise",
+    "set_exploration_modules_spec_from_env",
+    "set_recurrent_mode",
+    "RandomPolicy",
+]

torchrl/modules/distributions/__init__.py

@@ -0,0 +1,73 @@
+# 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 tensordict.nn import NormalParamExtractor
+from torch import distributions as torch_dist
+
+from .continuous import (
+    Delta,
+    IndependentNormal,
+    NormalParamWrapper,
+    TanhDelta,
+    TanhNormal,
+    TruncatedNormal,
+)
+from .discrete import (
+    LLMMaskedCategorical,
+    MaskedCategorical,
+    MaskedOneHotCategorical,
+    OneHotCategorical,
+    OneHotOrdinal,
+    Ordinal,
+    ReparamGradientStrategy,
+)
+
+distributions_maps = {
+    str(dist).lower(): dist
+    for dist in (
+        Delta,
+        IndependentNormal,
+        TanhDelta,
+        TanhNormal,
+        TruncatedNormal,
+        MaskedCategorical,
+        MaskedOneHotCategorical,
+        OneHotCategorical,
+        Ordinal,
+        OneHotOrdinal,
+    )
+}
+
+HAS_ENTROPY = {
+    Delta: False,
+    IndependentNormal: True,
+    TanhDelta: False,
+    TanhNormal: False,
+    TruncatedNormal: False,
+    MaskedCategorical: True,
+    MaskedOneHotCategorical: False,
+    OneHotCategorical: True,
+    torch_dist.Categorical: True,
+    torch_dist.Normal: True,
+}
+
+__all__ = [
+    "NormalParamExtractor",
+    "distributions",
+    "Delta",
+    "IndependentNormal",
+    "LLMMaskedCategorical",
+    "NormalParamWrapper",
+    "TanhDelta",
+    "TanhNormal",
+    "TruncatedNormal",
+    "MaskedCategorical",
+    "MaskedOneHotCategorical",
+    "OneHotCategorical",
+    "OneHotOrdinal",
+    "LLMMaskedCategorical",
+    "Ordinal",
+    "ReparamGradientStrategy",
+]