[Feature] Auto-configure exploration module specs from environment in Collector

docs/source/reference/modules_exploration.rst

@@ -13,3 +13,12 @@ Exploration modules add noise to actions to enable exploration during training.
     ConsistentDropoutModule
     EGreedyModule
     OrnsteinUhlenbeckProcessModule
+
+Helpers
+-------
+
+.. autosummary::
+    :toctree: generated/
+    :template: rl_template_noinherit.rst
+
+    set_exploration_modules_spec_from_env

test/test_configs.py

@@ -1008,6 +1008,39 @@ def test_value_model_config(self):
         assert value_model.module.in_features == 10
         assert value_model.module.out_features == 1
 
+    @pytest.mark.skipif(not _has_hydra, reason="Hydra is not installed")
+    def test_additive_gaussian_module_config(self):
+        """Test AdditiveGaussianModuleConfig."""
+        from hydra.utils import instantiate
+        from torchrl.trainers.algorithms.configs.modules import (
+            AdditiveGaussianModuleConfig,
+        )
+
+        cfg = AdditiveGaussianModuleConfig(
+            spec=None,
+            sigma_init=1.0,
+            sigma_end=0.1,
+            annealing_num_steps=1000,
+            mean=0.0,
+            std=0.1,
+            action_key="action",
+        )
+        assert (
+            cfg._target_
+            == "torchrl.trainers.algorithms.configs.modules._make_additive_gaussian_module"
+        )
+        assert cfg.spec is None
+        assert cfg.sigma_init == 1.0
+        assert cfg.sigma_end == 0.1
+        assert cfg.action_key == "action"
+
+        module = instantiate(cfg)
+        from torchrl.modules.tensordict_module.exploration import AdditiveGaussianModule
+
+        assert isinstance(module, AdditiveGaussianModule)
+        assert module._spec is None
+        assert module.action_key == "action"
+
 
 @pytest.mark.skipif(
     not _python_version_compatible, reason="Python 3.10+ required for config system"
@@ -1096,9 +1129,100 @@ def test_collector_config(self, factory, collector):
                 # Just check that we can iterate
                 break
         finally:
-            # Only call shutdown if the collector has that method
-            if hasattr(collector_instance, "shutdown"):
-                collector_instance.shutdown(timeout=10)
+            collector_instance.shutdown(timeout=10)
+
+    @pytest.mark.parametrize("factory", [True, False])
+    @pytest.mark.parametrize("collector", ["async", "multi_sync", "multi_async"])
+    @pytest.mark.skipif(not _has_gym, reason="Gym is not installed")
+    @pytest.mark.skipif(not _has_hydra, reason="Hydra is not installed")
+    def test_collector_auto_configures_exploration_modules(self, factory, collector):
+        """Test that collector instantiation auto-configures exploration modules.
+
+        This tests that when an exploration module (e.g. AdditiveGaussianModule) is configured
+        without a spec (spec=None), the collector's __init__ automatically sets the spec
+        from the environment's action_spec.
+        """
+        from hydra.utils import instantiate
+        from torchrl.collectors import (
+            AsyncCollector,
+            MultiAsyncCollector,
+            MultiSyncCollector,
+        )
+        from torchrl.trainers.algorithms.configs.collectors import (
+            AsyncDataCollectorConfig,
+            MultiAsyncCollectorConfig,
+            MultiSyncCollectorConfig,
+        )
+        from torchrl.trainers.algorithms.configs.envs_libs import GymEnvConfig
+        from torchrl.trainers.algorithms.configs.modules import (
+            AdditiveGaussianModuleConfig,
+            MLPConfig,
+            TanhNormalModelConfig,
+            TensorDictSequentialConfig,
+        )
+        from torchrl.trainers.algorithms.configs.weight_update import (
+            RemoteModuleWeightUpdaterConfig,
+        )
+
+        env_cfg = GymEnvConfig(env_name="Pendulum-v1")
+
+        policy_cfg = TanhNormalModelConfig(
+            network=MLPConfig(in_features=3, out_features=2, depth=2, num_cells=32),
+            in_keys=["observation"],
+            out_keys=["action"],
+        )
+
+        exploration_cfg = AdditiveGaussianModuleConfig(
+            spec=None,  # Will be auto-set by collector from environment
+            sigma_init=0.5,
+            sigma_end=0.1,
+            annealing_num_steps=100,
+            action_key="action",
+        )
+
+        exploratory_policy_cfg = TensorDictSequentialConfig(
+            modules=[policy_cfg, exploration_cfg],
+        )
+
+        if collector == "async":
+            cfg_cls = AsyncDataCollectorConfig
+            expected_cls = AsyncCollector
+            kwargs = {"create_env_fn": env_cfg, "frames_per_batch": 10}
+        elif collector == "multi_sync":
+            cfg_cls = MultiSyncCollectorConfig
+            expected_cls = MultiSyncCollector
+            kwargs = {"create_env_fn": [env_cfg], "frames_per_batch": 10}
+        elif collector == "multi_async":
+            cfg_cls = MultiAsyncCollectorConfig
+            expected_cls = MultiAsyncCollector
+            kwargs = {"create_env_fn": [env_cfg], "frames_per_batch": 10}
+        else:
+            raise ValueError(f"Unknown collector type: {collector}")
+
+        if factory:
+            cfg = cfg_cls(
+                policy_factory=exploratory_policy_cfg,
+                weight_updater=RemoteModuleWeightUpdaterConfig(),
+                **kwargs,
+            )
+        else:
+            cfg = cfg_cls(policy=exploratory_policy_cfg, **kwargs)
+
+        collector_instance = instantiate(cfg)
+
+        try:
+            assert isinstance(collector_instance, expected_cls)
+            # This would raise RuntimeError if spec was not set on workers' exploration module
+            for batch in collector_instance:
+                assert "action" in batch.keys()
+                assert batch["action"].shape[-1] == 1
+                assert (batch["action"] >= -2.0).all() and (
+                    batch["action"] <= 2.0
+                ).all(), "Actions should be clipped to environment spec bounds [-2, 2]"
+                break
+
+        finally:
+            collector_instance.shutdown(timeout=10)
 
 
 @pytest.mark.skipif(

test/test_exploration.py

@@ -621,8 +621,93 @@ def test_collector(self, device, parallel_spec, interface, seed=0):
         return
 
     def test_no_spec_error(self, device):
-        with pytest.raises(RuntimeError, match="spec cannot be None."):
-            AdditiveGaussianModule(spec=None).to(device)
+        # Test that forward() raises error if spec is None and not set via setter
+        module = AdditiveGaussianModule(spec=None, device=device)
+        action = torch.randn(10, 6, device=device)
+        with pytest.raises(RuntimeError, match="spec has not been set"):
+            module._add_noise(action)
+
+    def test_delayed_spec_initialization(self, device):
+        """Test that spec can be set via property setter after initialization."""
+        torch.manual_seed(0)
+        d_act = 6
+        action_spec = Bounded(
+            -torch.ones(d_act, device=device),
+            torch.ones(d_act, device=device),
+            (d_act,),
+            device=device,
+        )
+
+        module = AdditiveGaussianModule(spec=None, device=device)
+        assert module._spec is None
+
+        # Set spec via property setter
+        module.spec = action_spec
+        assert module._spec is not None
+
+        # Verify it works with forward
+        action = torch.randn(100, d_act, device=device)
+        noisy_action = module._add_noise(action)
+        assert action_spec.is_in(noisy_action)
+
+
+@pytest.mark.parametrize("device", get_default_devices())
+@pytest.mark.parametrize("use_batched_env", [False, True])
+def test_set_exploration_modules_spec_from_env(device, use_batched_env):
+    """Test set_exploration_modules_spec_from_env helper configures exploration modules."""
+    from tensordict.nn import TensorDictSequential
+    from torchrl.modules.tensordict_module.exploration import (
+        set_exploration_modules_spec_from_env,
+    )
+
+    torch.manual_seed(0)
+
+    if use_batched_env:
+        env = SerialEnv(2, ContinuousActionVecMockEnv)
+        env = env.to(device)
+        expected_spec = env.action_spec_unbatched
+    else:
+        env = ContinuousActionVecMockEnv(device=device)
+        expected_spec = env.action_spec
+    env.reset()
+
+    d_obs = env.observation_spec["observation"].shape[-1]
+    d_act = expected_spec.shape[-1]
+
+    # Create a policy with exploration module that has spec=None
+    net = nn.Sequential(
+        nn.Linear(d_obs, 2 * d_act, device=device), NormalParamExtractor()
+    )
+    module = SafeModule(
+        net,
+        in_keys=["observation"],
+        out_keys=["loc", "scale"],
+    )
+    policy = ProbabilisticActor(
+        module=module,
+        in_keys=["loc", "scale"],
+        out_keys=["action"],
+        distribution_class=TanhNormal,
+        default_interaction_type=InteractionType.RANDOM,
+    ).to(device)
+    exploration_module = AdditiveGaussianModule(spec=None, device=device)
+    exploratory_policy = TensorDictSequential(policy, exploration_module)
+
+    assert exploration_module._spec is None
+
+    set_exploration_modules_spec_from_env(exploratory_policy, env)
+
+    # Verify spec is set after configuration and matches the environment's action_spec
+    assert exploration_module._spec is not None
+    if isinstance(exploration_module._spec, Composite):
+        assert exploration_module._spec[exploration_module.action_key] == expected_spec
+    else:
+        assert exploration_module._spec == expected_spec
+
+    td = env.reset()
+    result = exploratory_policy(td)
+    assert "action" in result.keys()
+    env.close()
 
 
 @pytest.mark.parametrize("state_dim", [7])

torchrl/collectors/_single.py

@@ -42,6 +42,9 @@
     set_exploration_type,
 )
 from torchrl.modules import RandomPolicy
+from torchrl.modules.tensordict_module.exploration import (
+    set_exploration_modules_spec_from_env,
+)
 from torchrl.weight_update import WeightSyncScheme
 from torchrl.weight_update.utils import _resolve_model
 
@@ -405,6 +408,9 @@ def __init__(
         # Set up policy and weights
         self._setup_policy_and_weights(policy)
 
+        # Configure exploration modules with action_spec from environment
+        set_exploration_modules_spec_from_env(self.policy, self.env)
+
         # Apply environment device
         self._apply_env_device()
 

torchrl/modules/__init__.py

@@ -92,7 +92,10 @@
     VmapModule,
     WorldModelWrapper,
 )
-from .tensordict_module.exploration import RandomPolicy
+from .tensordict_module.exploration import (
+    RandomPolicy,
+    set_exploration_modules_spec_from_env,
+)
 from .utils import get_primers_from_module
 from .planners import CEMPlanner, MPCPlannerBase, MPPIPlanner  # usort:skip
 
@@ -183,6 +186,7 @@
     "get_primers_from_module",
     "recurrent_mode",
     "reset_noise",
+    "set_exploration_modules_spec_from_env",
     "set_recurrent_mode",
     "RandomPolicy",
 ]