[Test] Add automated Config/Class kwarg-parity checks for Hydra configs

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"
 )

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)

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)