Update (base update)

[ghstack-poisoned]

Author: vmoens

.github/workflows/docs.yml

@@ -26,7 +26,7 @@ jobs:
   build-docs:
     strategy:
       matrix:
-        python_version: [ "3.9" ]
+        python_version: [ "3.12" ]
         cuda_arch_version: [ "12.8" ]
     uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main
     with:
@@ -60,7 +60,7 @@ jobs:
         bash ./miniconda.sh -b -f -p "${conda_dir}"
         eval "$(${conda_dir}/bin/conda shell.bash hook)"
         printf "* Creating a test environment\n"
-        conda create --prefix "${env_dir}" -y python=3.9
+        conda create --prefix "${env_dir}" -y python=3.12
         printf "* Activating\n"
         conda activate "${env_dir}"
 

docs/requirements.txt

@@ -16,7 +16,7 @@ sphinx_design
 torchvision
 dm_control
 mujoco<3.3.6
-gym[classic_control,accept-rom-license,ale-py,atari]
+gymnasium[classic_control,atari]
 pygame
 tqdm
 ipython

docs/source/reference/config.rst

@@ -507,7 +507,7 @@ Training and Optimization Configurations
     SparseAdamConfig
 
 Logging Configurations
-~~~~~~~~~~~~~~~~~~~~~
+~~~~~~~~~~~~~~~~~~~~~~
 
 .. currentmodule:: torchrl.trainers.algorithms.configs.logging
 

docs/source/reference/envs.rst

@@ -1123,7 +1123,6 @@ to be able to create this other composition:
     ExcludeTransform
     FiniteTensorDictCheck
     FlattenObservation
-    FlattenTensorDict
     FrameSkipTransform
     GrayScale
     Hash

docs/source/reference/llms.rst

@@ -118,9 +118,9 @@ Usage
 Adding Custom Templates
 ^^^^^^^^^^^^^^^^^^^^^^^
 
-You can add custom chat templates for new model families using the :func:`torchrl.data.llm.chat.add_chat_template` function.
+You can add custom chat templates for new model families using the :func:`torchrl.data.llm.add_chat_template` function.
 
-.. autofunction:: torchrl.data.llm.chat.add_chat_template
+.. autofunction:: torchrl.data.llm.add_chat_template
 
 Usage Examples
 ^^^^^^^^^^^^^^
@@ -130,7 +130,7 @@ Adding a Llama Template
 
 .. code-block:: python
 
-    >>> from torchrl.data.llm.chat import add_chat_template, History
+    >>> from torchrl.data.llm import add_chat_template, History
     >>> from transformers import AutoTokenizer
     >>> 
     >>> # Define the Llama chat template

docs/source/reference/utils.rst

@@ -1,7 +1,7 @@
 .. currentmodule:: torchrl
 
 torchrl._utils package
-====================
+======================
 
 Set of utility methods that are used internally by the library.
 

pyproject.toml

@@ -149,4 +149,4 @@ first_party_detection = false
 [project.entry-points."vllm.general_plugins"]
 # Ensure FP32 overrides are registered in all vLLM processes (main, workers, and
 # the registry subprocess) before resolving model classes.
-fp32_overrides = "torchrl.modules.llm.backends.vllm_plugin:register_fp32_overrides"
+fp32_overrides = "torchrl.modules.llm.backends.vllm.vllm_plugin:register_fp32_overrides"

test/_utils_internal.py

@@ -21,13 +21,7 @@
 from tensordict.nn import TensorDictModuleBase
 from torch import nn, vmap
 
-from torchrl._utils import (
-    implement_for,
-    logger,
-    logger as torchrl_logger,
-    RL_WARNINGS,
-    seed_generator,
-)
+from torchrl._utils import implement_for, logger, RL_WARNINGS, seed_generator
 from torchrl.data.utils import CloudpickleWrapper
 from torchrl.envs import MultiThreadedEnv, ObservationNorm
 from torchrl.envs.batched_envs import ParallelEnv, SerialEnv
@@ -230,7 +224,7 @@ def f_retry(*args, **kwargs):
                     return f(*args, **kwargs)
                 except ExceptionToCheck as e:
                     msg = "%s, Retrying in %d seconds..." % (str(e), mdelay)
-                    torchrl_logger.info(msg)
+                    logger.info(msg)
                     time.sleep(mdelay)
                     mtries -= 1
             try:

test/llm/test_vllm.py

@@ -40,7 +40,7 @@ class TestAsyncVLLMIntegration:
     @pytest.mark.slow
     def test_vllm_api_compatibility(self, sampling_params):
         """Test that AsyncVLLM supports the same inputs as vLLM.LLM.generate()."""
-        from torchrl.modules.llm.backends.vllm_async import AsyncVLLM
+        from torchrl.modules.llm.backends import AsyncVLLM
 
         # Create AsyncVLLM service
         service = AsyncVLLM.from_pretrained(
@@ -113,7 +113,7 @@ def test_vllm_api_compatibility(self, sampling_params):
     def test_weight_updates_with_transformer(self, sampling_params):
         """Test weight updates using vLLMUpdater with a real transformer model."""
         from torchrl.collectors.llm.weight_update.vllm import vLLMUpdater
-        from torchrl.modules.llm.backends.vllm_async import AsyncVLLM
+        from torchrl.modules.llm.backends import AsyncVLLM
         from torchrl.modules.llm.policies.transformers_wrapper import (
             TransformersWrapper,
         )

test/llm/test_wrapper.py

@@ -18,7 +18,7 @@
 from tensordict.utils import _zip_strict
 from torchrl.data.llm import History
 from torchrl.envs.llm.transforms.kl import KLComputation, RetrieveKL, RetrieveLogProb
-from torchrl.modules.llm.backends.vllm_async import AsyncVLLM
+from torchrl.modules.llm import AsyncVLLM
 from torchrl.modules.llm.policies.common import (
     _batching,
     ChatHistory,