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

[bsprenger]

Description

This PR enables the use of exploration modules (e.g., AdditiveGaussianModule) with the Hydra config system by supporting delayed spec initialization. When policies are instantiated via configs, the environment's action_spec is not available at construction time, causing exploration modules that require spec at init to fail. This change allows spec=None at construction and auto-configures it when the collector is created.

Motivation and Context

When using the trainer's config YAML framework, policies are instantiated by Hydra before the environment is created inside the collector. Exploration modules like AdditiveGaussianModule previously required a spec argument at construction, which wasn't possible in this workflow since the spec depends on the environment.

This PR does the following:

1. Modifies AdditiveGaussianModule to support delayed spec setting. With this change, you do not need to provide a spec at initialization, but you must provide it before calling forward().
2. Adds a config dataclass for AdditiveGaussianModule
3. Implement a utility to iterate through a policy and set the spec on any exploration modules
4. Auto-configure the policy's exploration modules with the env spec inside the collector (if they exist)

• I have raised an issue to propose this change (required for new features and bug fixes) -- See [Feature Request] Support delayed spec initialization for exploration modules in Hydra config-based workflow #3318

Types of changes

• New feature (non-breaking change which adds core functionality)

Checklist

Go over all the following points, and put an x in all the boxes that apply.
If you are unsure about any of these, don't hesitate to ask. We are here to help!

• I have read the CONTRIBUTION guide (required)
• My change requires a change to the documentation.
• I have updated the tests accordingly (required for a bug fix or a new feature).
• I have updated the documentation accordingly.

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/rl/3317

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❌ 4 New Failures, 1 Cancelled Job, 3 Pending, 1 Unrelated Failure

As of commit 52308fe with merge base 0a99aaf ():

NEW FAILURES - The following jobs have failed:

• Continuous Benchmark (PR) / CPU Pytest benchmark (gh)
  Workflow failed! Resource not accessible by integration
• Continuous Benchmark / GPU Pytest benchmark (gh)
  Process completed with exit code 139.
• Habitat Tests on Linux / tests (3.9, 12.8) / linux-job (gh)
  RuntimeError: Command docker exec -t af1b804707e6e656372de8fe7c3ae32c6a21782daad82c76584b08c204d4874e /exec failed with exit code 1
• Libs Tests on Linux / unittests-gym (3.9, 12.8) / linux-job (gh)
  RuntimeError: Command docker exec -t 3e52a9b8ca6f3624a3379495f6b3577db2f75d99793ec6a8e2ed786429da938f /exec failed with exit code 1

CANCELLED JOB - The following job was cancelled. Please retry:

• Continuous Benchmark (PR) / GPU Pytest benchmark (gh)
  ##[error]The operation was canceled.

BROKEN TRUNK - The following job failed but were present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

• Unit-tests on Windows / unittests-cpu (3.10, windows.4xlarge, cpu) / windows-job (gh) (trunk failure)
  test/compile/test_compile_value.py::TestValueFunctionCompile::test_gae_with_tensor_gamma_compiles_fullgraph

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[vmoens]

LGTM! Thanks so much for this.

Do you think we need similar edits for OU and EGreedy? I think Claude can do that easily.

I made the following changes if they're ok with you:

1. torchrl/modules/tensordict_module/exploration.py

a) Updated AdditiveGaussianModule docstring (lines 255-262):

• Changed spec (TensorSpec) to spec (TensorSpec, optional)
• Added explanation that it can be None for delayed initialization
• Documented that spec must be set via property setter before calling forward()

b) Added example to set_exploration_modules_spec_from_env (lines 822-840):

• Added a complete doctest example showing how to use the helper function

2. torchrl/modules/__init__.py

a) Added import (line 95):

from .tensordict_module.exploration import RandomPolicy, set_exploration_modules_spec_from_env

b) Added to __all__ (line 186):

"set_exploration_modules_spec_from_env",

3. docs/source/reference/modules_exploration.rst

Added new "Helpers" section (lines 17-24):

Helpers
-------

.. autosummary::
    :toctree: generated/
    :template: rl_template_noinherit.rst

    set_exploration_modules_spec_from_env

---

4. test/test_configs.py

Removed unnecessary hasattr checks (2 locations):

• Line ~1131-1134: Removed defensive if hasattr(collector_instance, "shutdown"): check
• Line ~1226-1228: Removed defensive if hasattr(collector_instance, "shutdown"): check

Since shutdown is an abstract method in the base DataCollector class, these guards were unnecessary.