feat(policies): add a permutation-equivariant deep-set policy

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

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",
+]

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

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