torchrl-nightly 2025.8.8__cp312-cp312-manylinux1_x86_64.whl → 2025.8.10__cp312-cp312-manylinux1_x86_64.whl

benchmarks/ecosystem/vmas_rllib_vs_torchrl_sampling_performance.py

@@ -9,7 +9,6 @@ import pickle
 
 import time
 from pathlib import Path
-from typing import Dict
 
 import numpy as np
 
@@ -93,7 +92,7 @@ def run_vmas_rllib(
                 - result["timers"]["learn_time_ms"]
             )
 
-    def env_creator(config: Dict):
+    def env_creator(config: dict):
         env = vmas.make_env(
             scenario=config["scenario_name"],
             num_envs=config["num_envs"],

sota-implementations/grpo/grpo_utils.py

@@ -5,7 +5,8 @@
 from __future__ import annotations
 
 import time
-from typing import Any, Callable, Literal
+from collections.abc import Callable
+from typing import Any, Literal
 
 import torch
 from omegaconf import DictConfig

sota-implementations/redq/utils.py

@@ -4,8 +4,9 @@
 # LICENSE file in the root directory of this source tree.
 from __future__ import annotations
 
+from collections.abc import Callable, Sequence
+
 from copy import copy
-from typing import Callable, Sequence
 
 import torch
 from omegaconf import OmegaConf

torchrl/_utils.py

@@ -16,12 +16,13 @@ import threading
 import time
 import traceback
 import warnings
+from collections.abc import Callable
 from contextlib import nullcontext
 from copy import copy
 from functools import wraps
 from importlib import import_module
 from textwrap import indent
-from typing import Any, Callable, cast, TypeVar
+from typing import Any, cast, TypeVar
 
 import numpy as np
 import torch

torchrl/collectors/collectors.py

@@ -17,12 +17,13 @@ import time
 import typing
 import warnings
 from collections import defaultdict, OrderedDict
+from collections.abc import Callable, Iterator, Mapping, Sequence
 from copy import deepcopy
 from multiprocessing import connection, queues
 from multiprocessing.managers import SyncManager
 from queue import Empty
 from textwrap import indent
-from typing import Any, Callable, Iterator, Mapping, Sequence, TypeVar
+from typing import Any, TypeVar
 
 import numpy as np
 import torch

torchrl/collectors/distributed/generic.py

@@ -9,9 +9,11 @@ from __future__ import annotations
 import os
 import socket
 import warnings
+from collections import OrderedDict
+from collections.abc import Callable, Sequence
 from copy import copy, deepcopy
 from datetime import timedelta
-from typing import Any, Callable, OrderedDict, Sequence
+from typing import Any
 
 import torch.cuda
 from tensordict import TensorDict, TensorDictBase

torchrl/collectors/distributed/ray.py

@@ -7,7 +7,9 @@ from __future__ import annotations
 
 import asyncio
 import warnings
-from typing import Any, Callable, Iterator, OrderedDict, Sequence
+from collections import OrderedDict
+from collections.abc import Callable, Iterator, Sequence
+from typing import Any
 
 import torch
 import torch.nn as nn

torchrl/collectors/distributed/rpc.py

@@ -11,8 +11,10 @@ import os
 import socket
 import time
 import warnings
+from collections import OrderedDict
+from collections.abc import Callable, Sequence
 from copy import copy, deepcopy
-from typing import Any, Callable, OrderedDict, Sequence
+from typing import Any
 
 import torch.cuda
 

torchrl/collectors/distributed/sync.py

@@ -9,9 +9,11 @@ from __future__ import annotations
 import os
 import socket
 import warnings
+from collections import OrderedDict
+from collections.abc import Callable, Sequence
 from copy import copy, deepcopy
 from datetime import timedelta
-from typing import Any, Callable, Literal, OrderedDict, Sequence
+from typing import Any, Literal
 
 import torch.cuda
 from tensordict import TensorDict, TensorDictBase

torchrl/collectors/llm/base.py

@@ -5,7 +5,8 @@
 from __future__ import annotations
 
 from collections import deque
-from typing import Any, Callable
+from collections.abc import Callable
+from typing import Any
 
 import torch
 

torchrl/collectors/llm/ray_collector.py

@@ -7,7 +7,8 @@ from __future__ import annotations
 import copy
 
 import warnings
-from typing import Any, Callable, Iterator
+from collections.abc import Callable, Iterator
+from typing import Any
 
 import torch
 from tensordict import TensorDictBase

torchrl/collectors/utils.py

@@ -4,7 +4,7 @@
 # LICENSE file in the root directory of this source tree.
 from __future__ import annotations
 
-from typing import Callable
+from collections.abc import Callable
 
 import torch
 

torchrl/collectors/weight_update.py

@@ -6,7 +6,8 @@ from __future__ import annotations
 
 import abc
 import weakref
-from typing import Any, Callable, TypeVar
+from collections.abc import Callable
+from typing import Any, TypeVar
 
 import torch
 from tensordict import TensorDict, TensorDictBase

torchrl/data/datasets/atari_dqn.py

@@ -13,8 +13,8 @@ import shutil
 import subprocess
 import tempfile
 from collections import defaultdict
+from collections.abc import Callable
 from pathlib import Path
-from typing import Callable
 
 import numpy as np
 import torch

torchrl/data/datasets/common.py

@@ -6,8 +6,8 @@ from __future__ import annotations
 
 import abc
 import shutil
+from collections.abc import Callable
 from pathlib import Path
-from typing import Callable
 
 import torch
 from tensordict import TensorDictBase

torchrl/data/datasets/d4rl.py

@@ -10,8 +10,8 @@ import shutil
 import tempfile
 import urllib
 import warnings
+from collections.abc import Callable
 from pathlib import Path
-from typing import Callable
 
 import numpy as np
 import torch

torchrl/data/datasets/minari_data.py

@@ -10,10 +10,10 @@ import os.path
 import shutil
 import tempfile
 from collections import defaultdict
+from collections.abc import Callable
 from contextlib import nullcontext
 from dataclasses import asdict
 from pathlib import Path
-from typing import Callable
 
 import torch
 from tensordict import is_non_tensor, PersistentTensorDict, TensorDict

torchrl/data/datasets/openml.py

@@ -5,8 +5,8 @@
 from __future__ import annotations
 
 import os
+from collections.abc import Callable
 from pathlib import Path
-from typing import Callable
 
 import numpy as np
 from tensordict import TensorDict

torchrl/data/datasets/openx.py

@@ -10,8 +10,9 @@ import json
 import os
 import shutil
 import tempfile
+from collections.abc import Callable
 from pathlib import Path
-from typing import Any, Callable
+from typing import Any
 
 import torch
 from tensordict import make_tensordict, NonTensorData, pad, TensorDict

torchrl/data/datasets/roboset.py

@@ -8,9 +8,9 @@ import importlib.util
 import os.path
 import shutil
 import tempfile
+from collections.abc import Callable
 from contextlib import nullcontext
 from pathlib import Path
-from typing import Callable
 
 import torch
 from tensordict import PersistentTensorDict, TensorDict

torchrl/data/datasets/vd4rl.py

@@ -12,8 +12,8 @@ import pathlib
 import shutil
 import tempfile
 from collections import defaultdict
+from collections.abc import Callable
 from pathlib import Path
-from typing import Callable
 
 import numpy as np
 import torch

torchrl/data/llm/dataset.py

@@ -6,8 +6,8 @@ from __future__ import annotations
 
 import importlib.util
 import os
+from collections.abc import Sequence
 from pathlib import Path
-from typing import Sequence
 
 import torch
 from tensordict import TensorDict, TensorDictBase

torchrl/data/map/hash.py

@@ -4,7 +4,7 @@
 # LICENSE file in the root directory of this source tree.
 from __future__ import annotations
 
-from typing import Callable
+from collections.abc import Callable
 
 import torch
 from torch.nn import Module

torchrl/data/map/query.py

@@ -4,8 +4,10 @@
 # LICENSE file in the root directory of this source tree.
 from __future__ import annotations
 
+from collections.abc import Callable, Mapping
+
 from copy import deepcopy
-from typing import Any, Callable, List, Mapping, TypeVar
+from typing import Any, TypeVar
 
 import torch
 import torch.nn as nn
@@ -122,7 +124,7 @@ class QueryModule(TensorDictModuleBase):
     ):
         if len(in_keys) == 0:
             raise ValueError("`in_keys` cannot be empty.")
-        in_keys = in_keys if isinstance(in_keys, List) else [in_keys]
+        in_keys = in_keys if isinstance(in_keys, list) else [in_keys]
 
         super().__init__()
         in_keys = self.in_keys = in_keys

torchrl/data/map/tdstorage.py

@@ -7,7 +7,8 @@ from __future__ import annotations
 import abc
 import functools
 from abc import abstractmethod
-from typing import Any, Callable, Generic, TypeVar
+from collections.abc import Callable
+from typing import Any, Generic, TypeVar
 
 import torch
 from tensordict import is_tensor_collection, NestedKey, TensorDictBase

torchrl/data/map/tree.py

@@ -6,7 +6,8 @@ from __future__ import annotations
 
 import weakref
 from collections import deque
-from typing import Any, Callable, Literal
+from collections.abc import Callable
+from typing import Any, Literal
 
 import torch
 from tensordict import (

torchrl/data/map/utils.py

@@ -4,7 +4,7 @@
 # LICENSE file in the root directory of this source tree.
 from __future__ import annotations
 
-from typing import Callable
+from collections.abc import Callable
 
 from tensordict import NestedKey
 

torchrl/data/replay_buffers/ray_buffer.py

@@ -6,8 +6,9 @@ from __future__ import annotations
 
 import contextlib
 import importlib
+from collections.abc import Callable, Iterator
 
-from typing import Any, Callable, Iterator
+from typing import Any
 
 import torch
 from torchrl._utils import logger as torchrl_logger

torchrl/data/replay_buffers/replay_buffers.py

@@ -11,9 +11,10 @@ import multiprocessing
 import textwrap
 import threading
 import warnings
+from collections.abc import Callable, Sequence
 from concurrent.futures import ThreadPoolExecutor
 from pathlib import Path
-from typing import Any, Callable, Sequence
+from typing import Any
 
 import numpy as np
 import torch

torchrl/data/replay_buffers/scheduler.py

@@ -5,7 +5,8 @@
 from __future__ import annotations
 
 from abc import ABC, abstractmethod
-from typing import Any, Callable
+from collections.abc import Callable
+from typing import Any
 
 import numpy as np
 import torch

torchrl/data/replay_buffers/storages.py

@@ -11,9 +11,10 @@ import sys
 import textwrap
 import warnings
 from collections import OrderedDict
+from collections.abc import Callable, Mapping, Sequence
 from copy import copy
 from multiprocessing.context import get_spawning_popen
-from typing import Any, Callable, Mapping, Sequence
+from typing import Any
 
 import numpy as np
 import tensordict

torchrl/data/replay_buffers/utils.py

@@ -11,8 +11,9 @@ import math
 import operator
 import os
 import typing
+from collections.abc import Callable
 from pathlib import Path
-from typing import Any, Callable, Union
+from typing import Any, Union
 
 import numpy as np
 import torch

torchrl/data/replay_buffers/writers.py

@@ -8,10 +8,11 @@ import heapq
 import json
 import textwrap
 from abc import ABC, abstractmethod
+from collections.abc import Sequence
 from copy import copy
 from multiprocessing.context import get_spawning_popen
 from pathlib import Path
-from typing import Any, Sequence
+from typing import Any
 
 import numpy as np
 import torch

torchrl/data/tensor_specs.py

@@ -12,23 +12,12 @@ import gc
 import math
 import warnings
 import weakref
-from collections.abc import Iterable
+from collections.abc import Callable, Iterable, Sequence
 from copy import deepcopy
 from dataclasses import dataclass, field
 from functools import wraps
 from textwrap import indent
-from typing import (
-    Any,
-    Callable,
-    Dict,
-    Generic,
-    List,
-    overload,
-    Sequence,
-    Tuple,
-    TypeVar,
-    Union,
-)
+from typing import Any, Generic, overload, TypeVar, Union
 
 import numpy as np
 
@@ -61,27 +50,27 @@ except ImportError:
 
 DEVICE_TYPING = Union[torch.device, str, int]
 
-INDEX_TYPING = Union[int, torch.Tensor, np.ndarray, slice, List]
+INDEX_TYPING = Union[int, torch.Tensor, np.ndarray, slice, list]
 
 SHAPE_INDEX_TYPING = Union[
     int,
     range,
-    List[int],
+    list[int],
     np.ndarray,
     slice,
     None,
     torch.Tensor,
     type(...),
-    Tuple[
+    tuple[
         int,
         range,
-        List[int],
+        list[int],
         np.ndarray,
         slice,
         None,
         torch.Tensor,
         type(...),
-        Tuple[Any],
+        tuple[Any],
     ],
 ]
 
@@ -6273,7 +6262,7 @@ class StackedComposite(_LazyStackedMixin[Composite], Composite):
     def update(self, dict) -> None:
         for key, item in dict.items():
             if key in self.keys() and isinstance(
-                item, (Dict, Composite, StackedComposite)
+                item, (dict, Composite, StackedComposite)
             ):
                 for spec, sub_item in zip(self._specs, item.unbind(self.dim)):
                     spec[key].update(sub_item)

torchrl/data/utils.py

@@ -6,7 +6,8 @@ from __future__ import annotations
 
 import functools
 import typing
-from typing import Any, Callable, List, Tuple, Union
+from collections.abc import Callable
+from typing import Any, Union
 
 import numpy as np
 import torch
@@ -45,7 +46,7 @@ if hasattr(typing, "get_args"):
 else:
     DEVICE_TYPING_ARGS = (torch.device, str, int)
 
-INDEX_TYPING = Union[None, int, slice, str, Tensor, List[Any], Tuple[Any, ...]]
+INDEX_TYPING = Union[None, int, slice, str, Tensor, list[Any], tuple[Any, ...]]
 
 
 ACTION_SPACE_MAP = {

torchrl/envs/async_envs.py

@@ -7,12 +7,13 @@ from __future__ import annotations
 import abc
 
 import multiprocessing
+from collections.abc import Callable, Sequence
 from concurrent.futures import as_completed, ThreadPoolExecutor
 
 # import queue
 from multiprocessing import Queue
 from queue import Empty
-from typing import Callable, Literal, Sequence
+from typing import Literal
 
 import torch
 from tensordict import (

torchrl/envs/batched_envs.py

@@ -11,11 +11,12 @@ import os
 import time
 import weakref
 from collections import OrderedDict
+from collections.abc import Callable, Sequence
 from copy import deepcopy
 from functools import wraps
 from multiprocessing import connection
 from multiprocessing.synchronize import Lock as MpLock
-from typing import Any, Callable, Sequence
+from typing import Any
 from warnings import warn
 
 import torch

torchrl/envs/common.py

@@ -9,9 +9,10 @@ import abc
 import re
 import warnings
 import weakref
+from collections.abc import Callable, Iterator
 from copy import deepcopy
 from functools import partial, wraps
-from typing import Any, Callable, Iterator
+from typing import Any
 
 import numpy as np
 import torch

torchrl/envs/custom/llm.py

@@ -4,7 +4,7 @@
 # LICENSE file in the root directory of this source tree.
 from __future__ import annotations
 
-from typing import Callable
+from collections.abc import Callable
 
 import torch
 

torchrl/envs/env_creator.py

@@ -6,8 +6,8 @@
 from __future__ import annotations
 
 from collections import OrderedDict
+from collections.abc import Callable
 from multiprocessing.sharedctypes import Synchronized
-from typing import Callable
 
 import torch
 from tensordict import TensorDictBase

torchrl/envs/gym_like.py

@@ -9,7 +9,8 @@ import abc
 import functools
 import re
 import warnings
-from typing import Any, Callable, Mapping, Sequence, TypeVar
+from collections.abc import Callable, Mapping, Sequence
+from typing import Any, TypeVar
 
 import numpy as np
 import torch

torchrl/envs/libs/dm_control.py

@@ -7,7 +7,7 @@ from __future__ import annotations
 import collections
 import importlib
 import os
-from typing import Any, Dict
+from typing import Any
 
 import numpy as np
 import torch
@@ -44,7 +44,7 @@ def _dmcontrol_to_torchrl_spec_transform(
 ) -> TensorSpec:
     import dm_env
 
-    if isinstance(spec, collections.OrderedDict) or isinstance(spec, Dict):
+    if isinstance(spec, collections.OrderedDict) or isinstance(spec, dict):
         spec = {
             k: _dmcontrol_to_torchrl_spec_transform(
                 item,

torchrl/envs/libs/gym.py

@@ -10,7 +10,6 @@ import importlib
 import warnings
 from copy import copy
 from types import ModuleType
-from typing import Dict
 from warnings import warn
 
 import numpy as np
@@ -510,7 +509,7 @@ def convert_sequence_spec(
     return out
 
 
-@register_gym_spec_conversion(Dict)
+@register_gym_spec_conversion(dict)
 def convert_dict_spec(
     spec,
     dtype=None,
@@ -765,7 +764,7 @@ def _is_from_pixels(env):
     gDict = gym_backend("spaces").dict.Dict
     Box = gym_backend("spaces").Box
 
-    if isinstance(observation_spec, (Dict,)):
+    if isinstance(observation_spec, (dict,)):
         if "pixels" in set(observation_spec.keys()):
             return True
     if isinstance(observation_spec, (gDict,)):

torchrl/envs/libs/meltingpot.py

@@ -5,7 +5,7 @@
 from __future__ import annotations
 
 import importlib
-from typing import Mapping, Sequence
+from collections.abc import Mapping, Sequence
 
 import torch
 from tensordict import TensorDict, TensorDictBase

torchrl/envs/libs/pettingzoo.py

@@ -7,7 +7,6 @@ from __future__ import annotations
 import copy
 import importlib
 import warnings
-from typing import Dict
 
 import numpy as np
 import packaging
@@ -807,7 +806,7 @@ class PettingZooWrapper(_EnvWrapper):
                 for index, agent in enumerate(agents):
                     agent_obs = observation_dict[agent]
                     agent_info = info_dict[agent]
-                    if isinstance(agent_obs, Dict) and "action_mask" in agent_obs:
+                    if isinstance(agent_obs, dict) and "action_mask" in agent_obs:
                         if agent in agents_acting:
                             group_mask[index] = torch.tensor(
                                 agent_obs["action_mask"],
@@ -815,7 +814,7 @@ class PettingZooWrapper(_EnvWrapper):
                                 dtype=torch.bool,
                             )
                         del agent_obs["action_mask"]
-                    elif isinstance(agent_info, Dict) and "action_mask" in agent_info:
+                    elif isinstance(agent_info, dict) and "action_mask" in agent_info:
                         if agent in agents_acting:
                             group_mask[index] = torch.tensor(
                                 agent_info["action_mask"],