diff --git a/src/deephedging/__init__.py b/src/deephedging/__init__.py index f9ce6cf..2cb3d5d 100644 --- a/src/deephedging/__init__.py +++ b/src/deephedging/__init__.py @@ -76,6 +76,7 @@ ) from deephedging.market.local_vol import LocalVolSimulator from deephedging.policies import ( + DeepSetPolicy, FeedForwardPolicy, HedgePolicy, NoTransactionBandPolicy, @@ -107,6 +108,7 @@ "CorrelatedGBMSimulator", "CostModel", "DeepBSDESolver", + "DeepSetPolicy", "DefaultFeatures", "DiscountGenerator", "Entropic", diff --git a/src/deephedging/policies/__init__.py b/src/deephedging/policies/__init__.py index c195f53..2971afc 100644 --- a/src/deephedging/policies/__init__.py +++ b/src/deephedging/policies/__init__.py @@ -2,7 +2,14 @@ from deephedging.policies.band import NoTransactionBandPolicy from deephedging.policies.base import HedgePolicy +from deephedging.policies.deep_set import DeepSetPolicy from deephedging.policies.ffn import FeedForwardPolicy from deephedging.policies.recurrent import RecurrentPolicy -__all__ = ["FeedForwardPolicy", "HedgePolicy", "NoTransactionBandPolicy", "RecurrentPolicy"] +__all__ = [ + "DeepSetPolicy", + "FeedForwardPolicy", + "HedgePolicy", + "NoTransactionBandPolicy", + "RecurrentPolicy", +] diff --git a/src/deephedging/policies/deep_set.py b/src/deephedging/policies/deep_set.py new file mode 100644 index 0000000..81b177e --- /dev/null +++ b/src/deephedging/policies/deep_set.py @@ -0,0 +1,100 @@ +"""Permutation-equivariant deep-set hedging policy.""" + +import torch +from torch import nn + +from deephedging.policies.base import HedgePolicy + +_TOKEN_FEATURES = 3 + + +class DeepSetPolicy(HedgePolicy): + """Permutation-equivariant multi-asset policy in the Deep Sets family. + + A feedforward policy over a flat multi-asset feature vector relearns the + hedge for every permutation of the asset ordering and grows its first + layer with the asset count. This policy shares one encoder across assets + and conditions each asset's position on a pooled, order-invariant summary + of the book, so the map is equivariant under asset relabelling and the + parameter count is independent of how many assets trade. + + It pairs with :class:`~deephedging.features.MultiAssetFeatures`, whose + vector lays out the per-asset log moneyness, the shared time to maturity, + and the per-asset positions. One token per asset is reconstructed from + that layout as its own log moneyness and position alongside the shared + time, the encoder maps every token to a latent, mean pooling forms the + order-invariant summary, and the head reads each asset's position from its + own latent concatenated with the summary. + + Attributes: + n_assets: Number of assets the policy hedges. + encoder: Shared per-asset encoder. + head: Readout from an asset latent and the pooled summary. + """ + + def __init__( + self, + n_assets: int, + hidden_sizes: tuple[int, ...] = (32, 32), + latent_size: int = 32, + ) -> None: + """Initialises the policy network. + + Args: + n_assets: Number of assets on the book. + hidden_sizes: Widths of the shared encoder's hidden layers. + latent_size: Width of the per-asset latent and the pooled summary. + + Raises: + ValueError: If ``n_assets`` is not positive. + """ + super().__init__() + if n_assets < 1: + raise ValueError(f"n_assets must be at least 1, got {n_assets}") + self.n_assets = n_assets + encoder_layers: list[nn.Module] = [] + width = _TOKEN_FEATURES + for size in hidden_sizes: + encoder_layers.append(nn.Linear(width, size)) + encoder_layers.append(nn.SiLU()) + width = size + encoder_layers.append(nn.Linear(width, latent_size)) + self.encoder = nn.Sequential(*encoder_layers) + self.head = nn.Sequential( + nn.Linear(2 * latent_size, latent_size), + nn.SiLU(), + nn.Linear(latent_size, 1), + ) + + def forward( + self, features: torch.Tensor, state: torch.Tensor | None = None + ) -> tuple[torch.Tensor, torch.Tensor | None]: + """Computes the per-asset hedge positions for one rebalancing date. + + Args: + features: Multi-asset features of shape ``(n_paths, 2 * n_assets + + 1)`` laid out as :class:`~deephedging.features.MultiAssetFeatures` + produces them. + state: Ignored; present for interface compatibility. + + Returns: + Tuple of the positions per path, shape ``(n_paths, n_assets)``, + and ``None``. + + Raises: + ValueError: If the feature width disagrees with ``n_assets``. + """ + expected = 2 * self.n_assets + 1 + if features.shape[-1] != expected: + raise ValueError( + f"features must have width {expected} for {self.n_assets} assets, " + f"got {features.shape[-1]}" + ) + log_moneyness = features[:, : self.n_assets] + tau = features[:, self.n_assets : self.n_assets + 1] + position = features[:, self.n_assets + 1 :] + tokens = torch.stack((log_moneyness, position, tau.expand(-1, self.n_assets)), dim=-1) + latent = self.encoder(tokens) + pooled = latent.mean(dim=1, keepdim=True).expand(-1, self.n_assets, -1) + combined = torch.cat((latent, pooled), dim=-1) + return self.head(combined).squeeze(-1), None diff --git a/tests/unit/test_deep_set_policy.py b/tests/unit/test_deep_set_policy.py new file mode 100644 index 0000000..090d739 --- /dev/null +++ b/tests/unit/test_deep_set_policy.py @@ -0,0 +1,96 @@ +"""Tests for the permutation-equivariant deep-set policy.""" + +import pytest +import torch + +from deephedging.policies import DeepSetPolicy + + +def _features(n_paths: int, n_assets: int, seed: int = 5) -> torch.Tensor: + generator = torch.Generator().manual_seed(seed) + log_moneyness = 0.1 * torch.randn(n_paths, n_assets, generator=generator) + tau = torch.rand(n_paths, 1, generator=generator) + position = torch.randn(n_paths, n_assets, generator=generator) + return torch.cat((log_moneyness, tau, position), dim=-1) + + +def test_output_shape_matches_asset_count() -> None: + policy = DeepSetPolicy(n_assets=4) + position, state = policy(_features(16, 4)) + assert position.shape == (16, 4) + assert state is None + + +def test_policy_is_permutation_equivariant() -> None: + torch.manual_seed(0) + n_assets = 5 + policy = DeepSetPolicy(n_assets=n_assets) + features = _features(32, n_assets) + perm = torch.tensor([2, 0, 4, 1, 3]) + log_moneyness = features[:, :n_assets] + tau = features[:, n_assets : n_assets + 1] + position = features[:, n_assets + 1 :] + permuted = torch.cat((log_moneyness[:, perm], tau, position[:, perm]), dim=-1) + with torch.no_grad(): + base = policy(features)[0] + permuted_out = policy(permuted)[0] + assert torch.allclose(permuted_out, base[:, perm], atol=1e-5) + + +def test_parameter_count_is_independent_of_asset_count() -> None: + small = sum(p.numel() for p in DeepSetPolicy(n_assets=2).parameters()) + large = sum(p.numel() for p in DeepSetPolicy(n_assets=50).parameters()) + assert small == large + + +def test_rejects_bad_asset_count_and_feature_width() -> None: + with pytest.raises(ValueError): + DeepSetPolicy(n_assets=0) + policy = DeepSetPolicy(n_assets=3) + with pytest.raises(ValueError): + policy(torch.zeros(8, 6)) + + +@pytest.mark.slow +def test_deep_set_policy_trains_a_multi_asset_hedge() -> None: + from deephedging import CVaR, MultiAssetFeatures, TrainConfig, train + from deephedging.evaluation import expected_shortfall + from deephedging.frictions import NoCost + from deephedging.instruments import GeometricBasketCall + from deephedging.market import CorrelatedGBMSimulator, NoiseSpec + from deephedging.pricing import MonteCarloPricer + from deephedging.training import hedge_pnl + + torch.manual_seed(67) + sim = CorrelatedGBMSimulator( + s0=100.0, + sigmas=(0.2, 0.3), + correlation=((1.0, 0.5), (0.5, 1.0)), + maturity=0.25, + n_steps=8, + ) + payoff = GeometricBasketCall(strike=100.0) + premium = MonteCarloPricer(n_paths=200_000, seed=71).price(payoff, sim).value + feature_map = MultiAssetFeatures(n_assets=2) + policy = DeepSetPolicy(n_assets=2, hidden_sizes=(16, 16), latent_size=16) + config = TrainConfig(n_iterations=250, batch_paths=1024, lr=2e-3, seed=9) + train( + sim, + policy, + payoff, + NoCost(), + CVaR(alpha=0.9), + config, + premium=premium, + feature_map=feature_map, + ) + + eval_state = sim.simulate(50_000, noise=NoiseSpec(seed=73)) + with torch.no_grad(): + hedged = hedge_pnl( + eval_state, policy, payoff, NoCost(), premium=premium, feature_map=feature_map + ) + unhedged = premium - payoff(eval_state.spot) + assert float(expected_shortfall(hedged, alpha=0.9)) < 0.6 * float( + expected_shortfall(unhedged, alpha=0.9) + )