Update

[ghstack-poisoned]

Author: Vincent Moens

.github/workflows/lint.yml

@@ -35,7 +35,7 @@ jobs:
         echo '::endgroup::'
         
         echo '::group::Install lint tools'
-        pip install --progress-bar=off pre-commit
+        pip install --progress-bar=off pre-commit autoflake
         echo '::endgroup::'
         
         echo '::group::Lint Python source and configs'

build_tools/setup_helpers/__init__.py

@@ -3,6 +3,6 @@
 # This source code is licensed under the MIT license found in the
 # LICENSE file in the root directory of this source tree.
 
-from .extension import CMakeBuild, get_ext_modules  # noqa
+from .extension import CMakeBuild, get_ext_modules
 
 __all__ = ["CMakeBuild", "get_ext_modules"]

check_future_imports.py

@@ -1220,7 +1220,7 @@ Recorders are transforms that register data as they come in, for logging purpose
 
 Helpers
 -------
-.. currentmodule:: torchrl.envs.utils
+.. currentmodule:: torchrl.envs
 
 .. autosummary::
     :toctree: generated/

docs/source/reference/envs.rst

@@ -111,6 +111,7 @@ auto-completion to make their choice.
     :template: rl_template_noinherit.rst
 
     LossModule
+    add_random_module
 
 DQN
 ---

docs/source/reference/objectives.rst

@@ -100,7 +100,7 @@
     TD3BCLoss,
     TD3Loss,
 )
-from torchrl.objectives.common import LossModule
+from torchrl.objectives.common import add_random_module, LossModule
 from torchrl.objectives.deprecated import DoubleREDQLoss_deprecated, REDQLoss_deprecated
 from torchrl.objectives.redq import REDQLoss
 from torchrl.objectives.reinforce import ReinforceLoss
@@ -16162,6 +16162,15 @@ def _composite_log_prob(self):
         yield
         setter.unset()
 
+    def test_add_random_module(self):
+        class MyMod(nn.Module):
+            ...
+
+        add_random_module(MyMod)
+        import torchrl.objectives.utils
+
+        assert MyMod in torchrl.objectives.utils.RANDOM_MODULE_LIST
+
     def test_standardization(self):
         t = torch.arange(3 * 4 * 5 * 6, dtype=torch.float32).view(3, 4, 5, 6)
         std_t0 = _standardize(t, exclude_dims=(1, 3))

test/test_cost.py

@@ -5,7 +5,7 @@
 
 from .batched_envs import ParallelEnv, SerialEnv
 from .common import EnvBase, EnvMetaData, make_tensordict
-from .custom import ChessEnv, LLMHashingEnv, PendulumEnv, TicTacToeEnv
+from .custom import ChessEnv, LLMEnv, 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 (
@@ -46,6 +46,8 @@
 from .transforms import (
     ActionDiscretizer,
     ActionMask,
+    as_nested_tensor,
+    as_padded_tensor,
     AutoResetEnv,
     AutoResetTransform,
     BatchSizeTransform,
@@ -58,6 +60,7 @@
     Compose,
     ConditionalSkip,
     Crop,
+    DataLoadingPrimer,
     DeviceCastTransform,
     DiscreteActionProjection,
     DoubleToFloat,
@@ -116,134 +119,144 @@
     check_marl_grouping,
     exploration_type,
     ExplorationType,
+    get_available_libraries,
     make_composite_from_td,
     MarlGroupMapType,
+    RandomPolicy,
     set_exploration_type,
     step_mdp,
+    terminated_or_truncated,
 )
 
 __all__ = [
-    "ParallelEnv",
-    "SerialEnv",
-    "EnvBase",
-    "EnvMetaData",
-    "make_tensordict",
-    "ChessEnv",
-    "LLMHashingEnv",
-    "PendulumEnv",
-    "TicTacToeEnv",
-    "env_creator",
-    "EnvCreator",
-    "get_env_metadata",
-    "default_info_dict_reader",
-    "GymLikeEnv",
-    "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",
-    "DreamerDecoder",
-    "DreamerEnv",
-    "ModelBasedEnvBase",
     "ActionDiscretizer",
     "ActionMask",
     "AutoResetEnv",
     "AutoResetTransform",
     "BatchSizeTransform",
     "BinarizeReward",
+    "BraxEnv",
+    "BraxWrapper",
     "BurnInTransform",
     "CatFrames",
     "CatTensors",
     "CenterCrop",
+    "ChessEnv",
     "ClipTransform",
     "Compose",
     "ConditionalSkip",
     "Crop",
+    "DMControlEnv",
+    "DMControlWrapper",
+    "DTypeCastTransform",
+    "DataLoadingPrimer",
     "DeviceCastTransform",
     "DiscreteActionProjection",
     "DoubleToFloat",
-    "DTypeCastTransform",
+    "DreamerDecoder",
+    "DreamerEnv",
     "EndOfLifeTransform",
+    "EnvBase",
+    "EnvCreator",
+    "EnvMetaData",
     "ExcludeTransform",
+    "ExplorationType",
     "FiniteTensorDictCheck",
     "FlattenObservation",
     "FrameSkipTransform",
     "GrayScale",
-    "gSDENoise",
+    "GymEnv",
+    "GymLikeEnv",
+    "GymWrapper",
+    "HabitatEnv",
     "Hash",
     "InitTracker",
+    "IsaacGymEnv",
+    "IsaacGymWrapper",
+    "JumanjiEnv",
+    "JumanjiWrapper",
     "KLRewardTransform",
+    "LLMEnv",
+    "LLMHashingEnv",
     "LineariseRewards",
+    "MOGymEnv",
+    "MOGymWrapper",
+    "MarlGroupMapType",
+    "MeltingpotEnv",
+    "MeltingpotWrapper",
+    "ModelBasedEnvBase",
     "MultiAction",
     "MultiStepTransform",
+    "MultiThreadedEnv",
+    "MultiThreadedEnvWrapper",
     "NoopResetEnv",
     "ObservationNorm",
     "ObservationTransform",
+    "OpenMLEnv",
+    "OpenSpielEnv",
+    "OpenSpielWrapper",
+    "ParallelEnv",
+    "PendulumEnv",
     "PermuteTransform",
+    "PettingZooEnv",
+    "PettingZooWrapper",
     "PinMemoryTransform",
     "R3MTransform",
     "RandomCropTensorDict",
+    "RandomPolicy",
     "RemoveEmptySpecs",
     "RenameTransform",
     "Resize",
     "Reward2GoTransform",
     "RewardClipping",
     "RewardScaling",
     "RewardSum",
+    "RoboHiveEnv",
+    "SMACv2Env",
+    "SMACv2Wrapper",
     "SelectTransform",
+    "SerialEnv",
     "SignTransform",
     "SqueezeTransform",
     "Stack",
     "StepCounter",
     "TargetReturn",
     "TensorDictPrimer",
+    "TicTacToeEnv",
     "TimeMaxPool",
     "Timer",
-    "Tokenizer",
     "ToTensorImage",
+    "Tokenizer",
     "TrajCounter",
     "Transform",
     "TransformedEnv",
     "UnaryTransform",
+    "UnityMLAgentsEnv",
+    "UnityMLAgentsWrapper",
     "UnsqueezeTransform",
     "VC1Transform",
-    "VecGymEnvTransform",
-    "VecNorm",
     "VIPRewardTransform",
     "VIPTransform",
+    "VecGymEnvTransform",
+    "VecNorm",
+    "VmasEnv",
+    "VmasWrapper",
+    "as_nested_tensor",
+    "as_padded_tensor",
     "check_env_specs",
     "check_marl_grouping",
+    "default_info_dict_reader",
+    "env_creator",
     "exploration_type",
-    "ExplorationType",
+    "gSDENoise",
+    "get_available_libraries",
+    "get_env_metadata",
+    "gym_backend",
     "make_composite_from_td",
-    "MarlGroupMapType",
+    "make_tensordict",
+    "register_gym_spec_conversion",
     "set_exploration_type",
+    "set_gym_backend",
     "step_mdp",
+    "terminated_or_truncated",
 ]

torchrl/envs/__init__.py

@@ -4,8 +4,8 @@
 # LICENSE file in the root directory of this source tree.
 
 from .chess import ChessEnv
-from .llm import LLMHashingEnv
+from .llm import LLMEnv, LLMHashingEnv
 from .pendulum import PendulumEnv
 from .tictactoeenv import TicTacToeEnv
 
-__all__ = ["ChessEnv", "LLMHashingEnv", "PendulumEnv", "TicTacToeEnv"]
+__all__ = ["ChessEnv", "LLMHashingEnv", "PendulumEnv", "TicTacToeEnv", "LLMEnv"]