diff --git a/test/test_configs.py b/test/test_configs.py index 6dda89a8e0c..780091dd927 100644 --- a/test/test_configs.py +++ b/test/test_configs.py @@ -6,18 +6,21 @@ from __future__ import annotations import argparse +import dataclasses import importlib.util +import inspect import os +import pkgutil import subprocess import sys +import typing +import warnings import pytest import torch from tensordict.nn import TensorDictModule, TensorDictSequential - from torchrl import logger as torchrl_logger, trainers as trainers_module from torchrl.collectors import AsyncCollector, MultiAsyncCollector, MultiSyncCollector - from torchrl.data.replay_buffers.replay_buffers import ( ReplayBuffer, TensorDictReplayBuffer, @@ -85,6 +88,244 @@ pytestmark = pytest.mark.filterwarnings("error") +def _discover_leaf_configs() -> dict[str, type]: + """Collects every concrete ``*Config`` dataclass defined under the configs package. + + A class is "concrete" here if it declares its own ``_target_`` field (as + opposed to an abstract base like ``TransformConfig`` that subclasses share). + """ + from torchrl.trainers.algorithms.configs.common import ConfigBase + + configs = {} + for modinfo in pkgutil.iter_modules( + algorithm_configs.__path__, algorithm_configs.__name__ + "." + ): + module = importlib.import_module(modinfo.name) + for name, obj in vars(module).items(): + if ( + isinstance(obj, type) + and dataclasses.is_dataclass(obj) + and issubclass(obj, ConfigBase) + and obj.__module__ == module.__name__ + ): + configs[name] = obj + return configs + + +def _resolve_wrapped_class(target_path: str) -> type | None: + """Resolves a ``_target_`` dotted path to the class it ultimately constructs. + + ``_target_`` either names a class directly (most transforms, storages, + samplers, ...) or a ``_make_*``/``make_*`` factory function (trainers and a + few composite modules). For the factory case, the wrapped class is read off + the factory's return-type annotation rather than parsed out of its body. + Returns ``None`` if neither resolution path yields a class. + """ + module_path, _, attr = target_path.rpartition(".") + target = getattr(importlib.import_module(module_path), attr) + if inspect.isclass(target): + return target + if inspect.isfunction(target): + return_type = typing.get_type_hints(target).get("return") + if inspect.isclass(return_type): + return return_type + return None + + +_CONFIG_PARITY_UNRESOLVED = { + "BatchedEnvConfig": "make_batched_env dispatches to ParallelEnv, SerialEnv or " + "AsyncEnvPool based on batched_env_type, and per-backend kwargs pass through " + "the create_env_kwargs dict rather than individual fields; there is no single " + "wrapped class to diff kwargs against.", + "TanhNormalModelConfig": "_make_tanh_normal_model composes a TensorDictModule " + "and a ProbabilisticTensorDictModule into a ProbabilisticTensorDictSequential; " + "there is no single wrapped class whose __init__ the Config fields mirror.", + "StorageEnsembleWriterConfig": "_target_ references torchrl.data.replay_buffers." + "StorageEnsembleWriter, which does not exist in the current codebase.", + "KLRewardTransformConfig": "_target_ references torchrl.envs.transforms.llm, " + "which does not exist; KLRewardTransform now lives in " + "torchrl.envs.llm.transforms.kl.", + "LionConfig": "_target_ references torch.optim.Lion, which is not available in " + "the torch versions TorchRL currently supports.", +} + +_CONFIG_PARITY_KNOWN_GAPS = frozenset( + { + "ActionMaskConfig", + "AggregatorConfig", + "BurnInTransformConfig", + "CatTensorsConfig", + "CenterCropConfig", + "CropConfig", + "DTypeCastTransformConfig", + "DeviceCastTransformConfig", + "DiscreteActionProjectionConfig", + "ExcludeTransformConfig", + "FlattenObservationConfig", + "FlattenTensorDictConfig", + "HashConfig", + "ImmutableDatasetWriterConfig", + "IsaacGymEnvConfig", + "LayerConfig", + "LazyStackStorageConfig", + "MeltingpotEnvConfig", + "ModuleTransformConfig", + "MultiStepTransformConfig", + "MultiSyncCollectorConfig", + "MultiThreadedEnvConfig", + "NormConfig", + "ObservationNormConfig", + "OpenMLEnvConfig", + "OpenSpielEnvConfig", + "PermuteTransformConfig", + "PettingZooEnvConfig", + "PrioritizedSamplerConfig", + "PrioritizedSliceSamplerConfig", + "QValueModelConfig", + "R3MTransformConfig", + "RandomCropTensorDictConfig", + "RemoveEmptySpecsConfig", + "RenameTransformConfig", + "ReplayBufferConfig", + "Reward2GoTransformConfig", + "RewardSumConfig", + "SMACv2EnvConfig", + "SamplerEnsembleConfig", + "SelectTransformConfig", + "SignTransformConfig", + "SliceSamplerConfig", + "SliceSamplerWithoutReplacementConfig", + "SqueezeTransformConfig", + "StackConfig", + "StorageConfig", + "TensorDictModuleConfig", + "TensorDictPrimerConfig", + "TimeMaxPoolConfig", + "TokenizerConfig", + "UnaryTransformConfig", + "UnityMLAgentsEnvConfig", + "UnsqueezeTransformConfig", + "VC1TransformConfig", + "VIPRewardTransformConfig", + "VIPTransformConfig", + "ValueModelConfig", + "VecGymEnvTransformConfig", + "VecNormConfig", + "VecNormV2Config", + "WriterConfig", + } +) + + +def _has_concrete_target(cfg_cls: type) -> bool: + target_field = next( + (f for f in dataclasses.fields(cfg_cls) if f.name == "_target_"), None + ) + return target_field is not None and target_field.default not in ( + dataclasses.MISSING, + None, + ) + + +def _config_parity_cases() -> list: + if not _configs_available: + return [] + cases = [] + for name, cfg_cls in sorted(_discover_leaf_configs().items()): + if not _has_concrete_target(cfg_cls): + continue + if name in _CONFIG_PARITY_UNRESOLVED: + cases.append( + pytest.param( + name, + marks=pytest.mark.skip(reason=_CONFIG_PARITY_UNRESOLVED[name]), + ) + ) + elif name in _CONFIG_PARITY_KNOWN_GAPS: + cases.append( + pytest.param( + name, + marks=pytest.mark.xfail( + reason="pre-existing config/class kwarg-parity gap; see " + "CLAUDE.md section 14 (run with -rx for the missing kwargs)", + strict=True, + ), + ) + ) + else: + cases.append(name) + return cases + + +@pytest.mark.skipif( + not _python_version_compatible, reason="Python 3.10+ required for config system" +) +@pytest.mark.skipif( + not _configs_available, reason="Config system requires hydra-core and omegaconf" +) +class TestConfigClassParity: + """Checks that every leaf ``*Config`` dataclass stays in sync with the class it wraps. + + CLAUDE.md section 14 requires every kwarg accepted by a wrapped class's + ``__init__`` to appear as a Config field with the same default; otherwise a + user setting that kwarg through Hydra has no effect and the failure is + silent. This walks every concrete Config dataclass reachable from + ``torchrl.trainers.algorithms.configs``, resolves the class or factory named + by its ``_target_``, and asserts every required constructor parameter has a + same-named Config field. + + A handful of configs cannot be checked this way (composite or polymorphic + factories with no single wrapped class, or a ``_target_`` that references a + class which no longer exists) and are skipped with an explicit reason + instead of silently passing. A further set of configs have a known, + pre-existing gap and are marked ``xfail(strict=True)``: fixing one of them + without removing it from ``_CONFIG_PARITY_KNOWN_GAPS`` turns CI red, which + is the point -- the allowlist should only shrink. + """ + + @pytest.mark.parametrize("config_name", _config_parity_cases()) + def test_wrapped_class_kwargs_have_config_fields(self, config_name): + cfg_cls = _discover_leaf_configs()[config_name] + fields = {f.name: f for f in dataclasses.fields(cfg_cls)} + target_path = fields["_target_"].default + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + wrapped_cls = _resolve_wrapped_class(target_path) + assert wrapped_cls is not None, ( + f"{config_name}._target_ = {target_path!r} could not be resolved to " + "a class (not a class itself, and not a function with a class " + "return annotation)." + ) + + params = [ + (pname, param) + for pname, param in inspect.signature( + wrapped_cls.__init__ + ).parameters.items() + if pname != "self" + ] + if ( + fields.get("_partial_") is not None + and fields["_partial_"].default is True + and params + ): + params = params[1:] + + missing = [ + pname + for pname, param in params + if param.kind + not in (inspect.Parameter.VAR_POSITIONAL, inspect.Parameter.VAR_KEYWORD) + and pname not in fields + ] + assert not missing, ( + f"{config_name} is missing Config field(s) for " + f"{wrapped_cls.__name__}.__init__ kwarg(s) {missing}: these can be " + f"set on {wrapped_cls.__name__} directly but are silently " + f"unreachable through Hydra. See CLAUDE.md section 14." + ) + + @pytest.mark.skipif( not _python_version_compatible, reason="Python 3.10+ required for config system" ) diff --git a/torchrl/trainers/algorithms/configs/envs_libs.py b/torchrl/trainers/algorithms/configs/envs_libs.py index 585db4dfc05..8d86630c564 100644 --- a/torchrl/trainers/algorithms/configs/envs_libs.py +++ b/torchrl/trainers/algorithms/configs/envs_libs.py @@ -9,7 +9,7 @@ from typing import Any from omegaconf import MISSING -from torchrl.envs.libs.gym import set_gym_backend +from torchrl.envs.libs.gym import GymEnv, set_gym_backend from torchrl.envs.transforms.transforms import DoubleToFloat from torchrl.trainers.algorithms.configs.common import ConfigBase @@ -57,7 +57,7 @@ def make_gym_env( from_pixels: bool = False, double_to_float: bool = False, **kwargs, -): +) -> GymEnv: """Create a Gym/Gymnasium environment. Args: @@ -69,8 +69,6 @@ def make_gym_env( Returns: The created environment instance. """ - from torchrl.envs.libs.gym import GymEnv - if backend is not None: with set_gym_backend(backend): env = GymEnv(env_name, from_pixels=from_pixels, **kwargs) diff --git a/torchrl/trainers/algorithms/configs/modules.py b/torchrl/trainers/algorithms/configs/modules.py index 477d83fb8a2..46120b364d2 100644 --- a/torchrl/trainers/algorithms/configs/modules.py +++ b/torchrl/trainers/algorithms/configs/modules.py @@ -10,9 +10,14 @@ from typing import Any import torch - from omegaconf import MISSING - +from tensordict.nn import TensorDictModule, TensorDictSequential +from torchrl.modules import ( + AdditiveGaussianModule, + QValueActor, + TanhModule, + ValueOperator, +) from torchrl.trainers.algorithms.configs.common import ( _normalize_hydra_key, _normalize_hydra_keys, @@ -420,10 +425,9 @@ def __post_init__(self) -> None: super().__post_init__() -def _make_tensordict_module(*args, **kwargs): +def _make_tensordict_module(*args, **kwargs) -> TensorDictModule: """Helper function to create a TensorDictModule.""" from hydra.utils import instantiate - from tensordict.nn import TensorDictModule module = kwargs.pop("module") shared = kwargs.pop("shared", False) @@ -448,11 +452,10 @@ def _make_tensordict_module(*args, **kwargs): return tensordict_module -def _make_tensordict_sequential(*args, **kwargs): +def _make_tensordict_sequential(*args, **kwargs) -> TensorDictSequential: """Helper function to create a TensorDictSequential.""" from hydra.utils import instantiate from omegaconf import DictConfig, ListConfig - from tensordict.nn import TensorDictSequential modules = kwargs.pop("modules") shared = kwargs.pop("shared", False) @@ -549,12 +552,10 @@ def _make_tanh_normal_model(*args, **kwargs): return result -def _make_value_model(*args, **kwargs): +def _make_value_model(*args, **kwargs) -> ValueOperator: """Helper function to create a ValueOperator with the given network.""" from hydra.utils import instantiate - from torchrl.modules import ValueOperator - network = kwargs.pop("network") shared = kwargs.pop("shared", False) @@ -578,10 +579,8 @@ def _make_value_model(*args, **kwargs): return value_operator -def _make_tanh_module(*args, **kwargs): +def _make_tanh_module(*args, **kwargs) -> TanhModule: """Helper function to create a TanhModule.""" - from torchrl.modules import TanhModule - kwargs.pop("shared", False) if "in_keys" in kwargs: @@ -592,10 +591,8 @@ def _make_tanh_module(*args, **kwargs): return TanhModule(**kwargs) -def _make_additive_gaussian_module(*args, **kwargs): +def _make_additive_gaussian_module(*args, **kwargs) -> AdditiveGaussianModule: """Helper function to create an AdditiveGaussianModule.""" - from torchrl.modules.tensordict_module.exploration import AdditiveGaussianModule - kwargs.pop("shared", False) kwargs.pop("in_keys", None) kwargs.pop("out_keys", None) @@ -625,12 +622,10 @@ def __post_init__(self) -> None: super().__post_init__() -def _make_qvalue_model(*args, **kwargs): +def _make_qvalue_model(*args, **kwargs) -> QValueActor: """Helper function to create a QValueActor with the given network.""" from hydra.utils import instantiate - from torchrl.modules import QValueActor - network = kwargs.pop("network") shared = kwargs.pop("shared", False) kwargs.pop("out_keys", None) diff --git a/torchrl/trainers/algorithms/configs/trainers.py b/torchrl/trainers/algorithms/configs/trainers.py index 601131cfafd..b7cc22c5181 100644 --- a/torchrl/trainers/algorithms/configs/trainers.py +++ b/torchrl/trainers/algorithms/configs/trainers.py @@ -10,7 +10,6 @@ import torch from tensordict.nn import TensorDictModuleBase - from torchrl.collectors import BaseCollector from torchrl.data import Categorical, Composite, OneHot from torchrl.objectives.common import LossModule @@ -752,7 +751,6 @@ def __post_init__(self) -> None: def _make_dqn_trainer(*args, **kwargs) -> DQNTrainer: from tensordict.nn import TensorDictSequential - from torchrl.modules import EGreedyModule from torchrl.trainers.trainers import Logger @@ -1416,9 +1414,8 @@ def __post_init__(self) -> None: super().__post_init__() -def _make_td3_trainer(*args, **kwargs): +def _make_td3_trainer(*args, **kwargs) -> TD3Trainer: from tensordict.nn import TensorDictSequential - from torchrl.objectives.utils import TargetNetUpdater from torchrl.trainers.algorithms.td3 import TD3OptimizationStepper from torchrl.trainers.trainers import Logger