diff --git a/docs/source/reference/modules_exploration.rst b/docs/source/reference/modules_exploration.rst index 01b25ef05cf..b8ff8a798b9 100644 --- a/docs/source/reference/modules_exploration.rst +++ b/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 diff --git a/test/test_configs.py b/test/test_configs.py index 575ebe1ca9e..e55856f21f4 100644 --- a/test/test_configs.py +++ b/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( diff --git a/test/test_exploration.py b/test/test_exploration.py index 377cc7c1ec6..64576305f36 100644 --- a/test/test_exploration.py +++ b/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]) diff --git a/torchrl/collectors/_single.py b/torchrl/collectors/_single.py index dbb9043f67a..b528417048b 100644 --- a/torchrl/collectors/_single.py +++ b/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() diff --git a/torchrl/modules/__init__.py b/torchrl/modules/__init__.py index dc8b213d492..196a43f8091 100644 --- a/torchrl/modules/__init__.py +++ b/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", ] diff --git a/torchrl/modules/tensordict_module/exploration.py b/torchrl/modules/tensordict_module/exploration.py index 4f8abaa225e..c99e34c8d79 100644 --- a/torchrl/modules/tensordict_module/exploration.py +++ b/torchrl/modules/tensordict_module/exploration.py @@ -5,6 +5,7 @@ from __future__ import annotations import warnings +from typing import TYPE_CHECKING import numpy as np import torch @@ -21,12 +22,16 @@ from torchrl.envs.utils import exploration_type, ExplorationType from torchrl.modules.tensordict_module.common import _forward_hook_safe_action +if TYPE_CHECKING: + from torchrl.envs import EnvBase + __all__ = [ "EGreedyWrapper", "EGreedyModule", "AdditiveGaussianModule", "OrnsteinUhlenbeckProcessModule", "OrnsteinUhlenbeckProcessWrapper", + "set_exploration_modules_spec_from_env", ] @@ -248,8 +253,13 @@ class AdditiveGaussianModule(TensorDictModuleBase): """Additive Gaussian PO module. Args: - spec (TensorSpec): the spec used for sampling actions. The sampled + spec (TensorSpec, optional): the spec used for sampling actions. The sampled action will be projected onto the valid action space once explored. + Can be ``None`` for delayed initialization, in which case the spec + must be set via the :attr:`spec` property setter before calling + :meth:`forward`. This is useful for config-based workflows where the + spec depends on the environment. + default: None sigma_init (scalar, optional): initial epsilon value. default: 1.0 sigma_end (scalar, optional): final epsilon value. @@ -257,9 +267,9 @@ class AdditiveGaussianModule(TensorDictModuleBase): annealing_num_steps (int, optional): number of steps it will take for sigma to reach the :obj:`sigma_end` value. default: 1000 - mean (:obj:`float`, optional): mean of each output element’s normal distribution. + mean (:obj:`float`, optional): mean of each output element's normal distribution. default: 0.0 - std (:obj:`float`, optional): standard deviation of each output element’s normal distribution. + std (:obj:`float`, optional): standard deviation of each output element's normal distribution. default: 1.0 Keyword Args: @@ -284,7 +294,7 @@ class AdditiveGaussianModule(TensorDictModuleBase): def __init__( self, - spec: TensorSpec, + spec: TensorSpec | None = None, sigma_init: float = 1.0, sigma_end: float = 0.1, annealing_num_steps: int = 1000, @@ -315,11 +325,11 @@ def __init__( "sigma", torch.tensor(sigma_init, dtype=torch.float32, device=device) ) + # spec can be None for delayed initialization + # In this case, spec must be set before forward() is called if spec is not None: if not isinstance(spec, Composite) and len(self.out_keys) >= 1: spec = Composite({action_key: spec}, shape=spec.shape[:-1]) - else: - raise RuntimeError("spec cannot be None.") self._spec = spec self.safe = safe if self.safe: @@ -329,6 +339,23 @@ def __init__( def spec(self): return self._spec + @spec.setter + def spec(self, value: TensorSpec) -> None: + """Set the action spec. + + This setter allows delayed initialization of the spec, which is useful + when the spec depends on environment information that is not available + at module construction time. + + Args: + value (TensorSpec): The action spec to set. Cannot be None. + """ + if value is None: + raise RuntimeError("spec cannot be set to None.") + if not isinstance(value, Composite) and len(self.out_keys) >= 1: + value = Composite({self.action_key: value}, shape=value.shape[:-1]) + self._spec = value + def step(self, frames: int = 1) -> None: """A step of sigma decay. @@ -350,6 +377,11 @@ def step(self, frames: int = 1) -> None: ) def _add_noise(self, action: torch.Tensor) -> torch.Tensor: + if self._spec is None: + raise RuntimeError( + "spec has not been set. The spec must be provided either at construction " + "time or set via the `spec` property before calling forward(). " + ) sigma = self.sigma mean = self.mean.expand(action.shape) std = self.std.expand(action.shape) @@ -771,3 +803,65 @@ def __call__(self, td: TensorDictBase) -> TensorDictBase: return td.update(self.action_spec.rand()) else: return td.set(self.action_key, self.action_spec.rand()) + + +def set_exploration_modules_spec_from_env( + policy: nn.Module | None, env: EnvBase +) -> None: + """Set spec on exploration modules in a policy with action_spec from environment. + + This helper function automatically sets the action_spec on exploration modules + (like AdditiveGaussianModule) that were instantiated without a spec. This + enables delayed initialization when using config systems where the spec + depends on the environment and should be set inside the collector. + + Args: + policy: The policy module (may contain exploration modules). Can be None. + env: The environment to extract action_spec from. + + Examples: + >>> import torch + >>> from tensordict.nn import TensorDictSequential + >>> from torchrl.envs import GymEnv + >>> from torchrl.modules import AdditiveGaussianModule, Actor, set_exploration_modules_spec_from_env + >>> # Create environment + >>> env = GymEnv("Pendulum-v1") + >>> # Create policy with exploration module without spec + >>> actor = Actor( + ... module=torch.nn.Linear(3, 1), + ... in_keys=["observation"], + ... out_keys=["action"], + ... ) + >>> exploration = AdditiveGaussianModule(spec=None) # spec will be set later + >>> policy = TensorDictSequential(actor, exploration) + >>> # Set spec from environment + >>> set_exploration_modules_spec_from_env(policy, env) + >>> assert exploration._spec is not None # spec is now set + >>> env.close() + + """ + if policy is None: + return + + # Only process nn.Module policies that have .modules() method + if not isinstance(policy, nn.Module): + return + + if hasattr(env, "action_spec_unbatched"): + action_spec = env.action_spec_unbatched + elif hasattr(env, "action_spec"): + action_spec = env.action_spec + else: + return + + if action_spec is None: + return + + # Types of exploration modules that support delayed spec initialization. + # More exploration modules can be added when they support delayed spec setting. + exploration_module_types = (AdditiveGaussianModule,) + + for submodule in policy.modules(): + if isinstance(submodule, exploration_module_types): + if hasattr(submodule, "_spec") and submodule._spec is None: + submodule.spec = action_spec diff --git a/torchrl/trainers/algorithms/configs/__init__.py b/torchrl/trainers/algorithms/configs/__init__.py index 5ff626f30cd..6bef247a6ab 100644 --- a/torchrl/trainers/algorithms/configs/__init__.py +++ b/torchrl/trainers/algorithms/configs/__init__.py @@ -81,6 +81,7 @@ WandbLoggerConfig, ) from torchrl.trainers.algorithms.configs.modules import ( + AdditiveGaussianModuleConfig, ConvNetConfig, MLPConfig, ModelConfig, @@ -270,6 +271,7 @@ "TensorDictModuleConfig", "TensorDictSequentialConfig", "ValueModelConfig", + "AdditiveGaussianModuleConfig", # Transforms - Core "ActionDiscretizerConfig", "ActionMaskConfig", @@ -447,6 +449,13 @@ def _register_configs(): cs.store(group="model", name="tanh_normal", node=TanhNormalModelConfig) cs.store(group="model", name="value", node=ValueModelConfig) + # Exploration configs + cs.store( + group="exploration", + name="additive_gaussian", + node=AdditiveGaussianModuleConfig, + ) + # ============================================================================= # Transform Configurations # ============================================================================= diff --git a/torchrl/trainers/algorithms/configs/modules.py b/torchrl/trainers/algorithms/configs/modules.py index a2477958c43..9eb6afa800d 100644 --- a/torchrl/trainers/algorithms/configs/modules.py +++ b/torchrl/trainers/algorithms/configs/modules.py @@ -347,6 +347,43 @@ def __post_init__(self) -> None: super().__post_init__() +@dataclass +class AdditiveGaussianModuleConfig(ModelConfig): + """A class to configure an AdditiveGaussianModule. + + Example: + >>> cfg = AdditiveGaussianModuleConfig( + ... spec=None, + ... sigma_init=1.0, + ... sigma_end=0.1, + ... mean=0.0, + ... std=1.0, + ... action_key="action", + ... ) + >>> module = instantiate(cfg) + >>> assert isinstance(module, AdditiveGaussianModule) + + .. seealso:: :class:`torchrl.modules.AdditiveGaussianModule` + """ + + spec: Any = None + sigma_init: float = 1.0 + sigma_end: float = 0.1 + annealing_num_steps: int = 1000 + mean: float = 0.0 + std: float = 1.0 + action_key: Any = "action" + safe: bool = False + device: Any = None + _target_: str = ( + "torchrl.trainers.algorithms.configs.modules._make_additive_gaussian_module" + ) + _partial_: bool = False + + def __post_init__(self) -> None: + super().__post_init__() + + def _make_tensordict_module(*args, **kwargs): """Helper function to create a TensorDictModule.""" from hydra.utils import instantiate @@ -511,3 +548,23 @@ def _make_tanh_module(*args, **kwargs): kwargs["out_keys"] = list(kwargs["out_keys"]) return TanhModule(**kwargs) + + +def _make_additive_gaussian_module(*args, **kwargs): + """Helper function to create an AdditiveGaussianModule.""" + from omegaconf import ListConfig + + from torchrl.modules.tensordict_module.exploration import AdditiveGaussianModule + + kwargs.pop("shared", False) + kwargs.pop("in_keys", None) + kwargs.pop("out_keys", None) + + if "action_key" in kwargs and isinstance(kwargs["action_key"], ListConfig): + action_key_list = list(kwargs["action_key"]) + if len(action_key_list) == 1: + kwargs["action_key"] = action_key_list[0] + else: + kwargs["action_key"] = tuple(action_key_list) + + return AdditiveGaussianModule(**kwargs)