Update

[ghstack-poisoned]

Author: Vincent Moens

torchrl/__init__.py

@@ -46,6 +46,20 @@
             )
 
 
+import torchrl.collectors
+import torchrl.data
+import torchrl.envs
+import torchrl.modules
+import torchrl.objectives
+import torchrl.trainers
+from torchrl._utils import (
+    auto_unwrap_transformed_env,
+    compile_with_warmup,
+    implement_for,
+    set_auto_unwrap_transformed_env,
+    timeit,
+)
+
 # Filter warnings in subprocesses: True by default given the multiple optional
 # deps of the library. This can be turned on via `torchrl.filter_warnings_subprocess = False`.
 filter_warnings_subprocess = True

torchrl/collectors/__init__.py

@@ -2,3 +2,13 @@
 #
 # This source code is licensed under the MIT license found in the
 # LICENSE file in the root directory of this source tree.
+
+from torchrl.envs.utils import RandomPolicy
+
+from .collectors import (
+    aSyncDataCollector,
+    DataCollectorBase,
+    MultiaSyncDataCollector,
+    MultiSyncDataCollector,
+    SyncDataCollector,
+)

torchrl/collectors/distributed/__init__.py

@@ -2,3 +2,9 @@
 #
 # This source code is licensed under the MIT license found in the
 # LICENSE file in the root directory of this source tree.
+
+from .generic import DEFAULT_SLURM_CONF, DistributedDataCollector
+from .ray import RayCollector
+from .rpc import RPCDataCollector
+from .sync import DistributedSyncDataCollector
+from .utils import submitit_delayed_launcher

torchrl/data/map/__init__.py

@@ -2,3 +2,8 @@
 #
 # This source code is licensed under the MIT license found in the
 # LICENSE file in the root directory of this source tree.
+
+from .hash import BinaryToDecimal, RandomProjectionHash, SipHash
+from .query import HashToInt, QueryModule
+from .tdstorage import TensorDictMap, TensorMap
+from .tree import MCTSForest, Tree

torchrl/data/postprocs/__init__.py

@@ -2,3 +2,5 @@
 #
 # This source code is licensed under the MIT license found in the
 # LICENSE file in the root directory of this source tree.
+
+from .postprocs import MultiStep

torchrl/data/replay_buffers/__init__.py

@@ -2,3 +2,49 @@
 #
 # This source code is licensed under the MIT license found in the
 # LICENSE file in the root directory of this source tree.
+
+from .checkpointers import (
+    FlatStorageCheckpointer,
+    H5StorageCheckpointer,
+    ListStorageCheckpointer,
+    NestedStorageCheckpointer,
+    StorageCheckpointerBase,
+    StorageEnsembleCheckpointer,
+    TensorStorageCheckpointer,
+)
+from .replay_buffers import (
+    PrioritizedReplayBuffer,
+    RemoteTensorDictReplayBuffer,
+    ReplayBuffer,
+    ReplayBufferEnsemble,
+    TensorDictPrioritizedReplayBuffer,
+    TensorDictReplayBuffer,
+)
+from .samplers import (
+    PrioritizedSampler,
+    PrioritizedSliceSampler,
+    RandomSampler,
+    Sampler,
+    SamplerEnsemble,
+    SamplerWithoutReplacement,
+    SliceSampler,
+    SliceSamplerWithoutReplacement,
+)
+from .storages import (
+    LazyMemmapStorage,
+    LazyStackStorage,
+    LazyTensorStorage,
+    ListStorage,
+    Storage,
+    StorageEnsemble,
+    TensorStorage,
+)
+from .utils import Flat2TED, H5Combine, H5Split, Nested2TED, TED2Flat, TED2Nested
+from .writers import (
+    ImmutableDatasetWriter,
+    RoundRobinWriter,
+    TensorDictMaxValueWriter,
+    TensorDictRoundRobinWriter,
+    Writer,
+    WriterEnsemble,
+)

torchrl/data/replay_buffers/storages.py

@@ -416,7 +416,7 @@ def __init__(
         max_size: int | None = None,
         *,
         compilable: bool = False,
-        stack_dim: int = -1,
+        stack_dim: int = 0,
     ):
         super().__init__(max_size=max_size, compilable=compilable)
         self.stack_dim = stack_dim

torchrl/data/rlhf/__init__.py

@@ -2,3 +2,13 @@
 #
 # This source code is licensed under the MIT license found in the
 # LICENSE file in the root directory of this source tree.
+
+from .dataset import (
+    create_infinite_iterator,
+    get_dataloader,
+    TensorDictTokenizer,
+    TokenizedDatasetLoader,
+)
+from .prompt import PromptData, PromptTensorDictTokenizer
+from .reward import PairwiseDataset, RewardData
+from .utils import AdaptiveKLController, ConstantKLController, RolloutFromModel

torchrl/envs/__init__.py

@@ -2,3 +2,122 @@
 #
 # This source code is licensed under the MIT license found in the
 # LICENSE file in the root directory of this source tree.
+
+from .batched_envs import ParallelEnv, SerialEnv
+from .common import EnvBase, EnvMetaData, make_tensordict
+from .custom import ChessEnv, LLMHashingEnv, PendulumEnv, TicTacToeEnv
+from .env_creator import env_creator, EnvCreator, get_env_metadata
+from .gym_like import default_info_dict_reader, GymLikeEnv
+from .libs import (
+    BraxEnv,
+    BraxWrapper,
+    DMControlEnv,
+    DMControlWrapper,
+    gym_backend,
+    GymEnv,
+    GymWrapper,
+    HabitatEnv,
+    IsaacGymEnv,
+    IsaacGymWrapper,
+    JumanjiEnv,
+    JumanjiWrapper,
+    MeltingpotEnv,
+    MeltingpotWrapper,
+    MOGymEnv,
+    MOGymWrapper,
+    MultiThreadedEnv,
+    MultiThreadedEnvWrapper,
+    OpenMLEnv,
+    OpenSpielEnv,
+    OpenSpielWrapper,
+    PettingZooEnv,
+    PettingZooWrapper,
+    register_gym_spec_conversion,
+    RoboHiveEnv,
+    set_gym_backend,
+    SMACv2Env,
+    SMACv2Wrapper,
+    UnityMLAgentsEnv,
+    UnityMLAgentsWrapper,
+    VmasEnv,
+    VmasWrapper,
+)
+from .model_based import DreamerDecoder, DreamerEnv, ModelBasedEnvBase
+from .transforms import (
+    ActionDiscretizer,
+    ActionMask,
+    AutoResetEnv,
+    AutoResetTransform,
+    BatchSizeTransform,
+    BinarizeReward,
+    BurnInTransform,
+    CatFrames,
+    CatTensors,
+    CenterCrop,
+    ClipTransform,
+    Compose,
+    ConditionalSkip,
+    Crop,
+    DeviceCastTransform,
+    DiscreteActionProjection,
+    DoubleToFloat,
+    DTypeCastTransform,
+    EndOfLifeTransform,
+    ExcludeTransform,
+    FiniteTensorDictCheck,
+    FlattenObservation,
+    FrameSkipTransform,
+    GrayScale,
+    gSDENoise,
+    Hash,
+    InitTracker,
+    KLRewardTransform,
+    LineariseRewards,
+    MultiAction,
+    MultiStepTransform,
+    NoopResetEnv,
+    ObservationNorm,
+    ObservationTransform,
+    PermuteTransform,
+    PinMemoryTransform,
+    R3MTransform,
+    RandomCropTensorDict,
+    RemoveEmptySpecs,
+    RenameTransform,
+    Resize,
+    Reward2GoTransform,
+    RewardClipping,
+    RewardScaling,
+    RewardSum,
+    SelectTransform,
+    SignTransform,
+    SqueezeTransform,
+    Stack,
+    StepCounter,
+    TargetReturn,
+    TensorDictPrimer,
+    TimeMaxPool,
+    Timer,
+    Tokenizer,
+    ToTensorImage,
+    TrajCounter,
+    Transform,
+    TransformedEnv,
+    UnaryTransform,
+    UnsqueezeTransform,
+    VC1Transform,
+    VecGymEnvTransform,
+    VecNorm,
+    VIPRewardTransform,
+    VIPTransform,
+)
+from .utils import (
+    check_env_specs,
+    check_marl_grouping,
+    exploration_type,
+    ExplorationType,
+    make_composite_from_td,
+    MarlGroupMapType,
+    set_exploration_type,
+    step_mdp,
+)

torchrl/envs/custom/__init__.py

@@ -2,3 +2,8 @@
 #
 # This source code is licensed under the MIT license found in the
 # LICENSE file in the root directory of this source tree.
+
+from .chess import ChessEnv
+from .llm import LLMHashingEnv
+from .pendulum import PendulumEnv
+from .tictactoeenv import TicTacToeEnv