diff --git a/src/yahtzee_agent/model.py b/src/yahtzee_agent/model.py index 7d3e265..ba8664c 100644 --- a/src/yahtzee_agent/model.py +++ b/src/yahtzee_agent/model.py @@ -11,7 +11,7 @@ from utilities.sequential_block import SequentialBlock from .features import PhiFeature -from .modules import Block, RollingHead, ScoringHead, ValueHead +from .modules import Block, ResidualBlock, RollingHead, ScoringHead, ValueHead class RollingActionRepresentation(str, Enum): @@ -173,9 +173,10 @@ def __init__( # noqa: PLR0913 dice_output_size = len(DICE_MASKS) # 32 possible masks scoring_output_size = 13 - layers = [Block(input_size, hidden_size, dropout_rate, activation)] + # Use ResidualBlock for trunk to prevent critic head from decohering shared representation + layers = [ResidualBlock(input_size, hidden_size, dropout_rate, activation)] for _ in range(num_hidden - 2): - layers.append(Block(hidden_size, hidden_size, dropout_rate, activation)) # noqa: PERF401 + layers.append(ResidualBlock(hidden_size, hidden_size, dropout_rate, activation)) # noqa: PERF401 self.network = SequentialBlock(*layers) diff --git a/src/yahtzee_agent/modules/__init__.py b/src/yahtzee_agent/modules/__init__.py index 37c0ce0..95ff39e 100644 --- a/src/yahtzee_agent/modules/__init__.py +++ b/src/yahtzee_agent/modules/__init__.py @@ -2,8 +2,9 @@ from .block import Block from .masked_softmax import MaskedSoftmax +from .residual_block import ResidualBlock from .rolling_head import RollingHead from .scoring_head import ScoringHead from .value_head import ValueHead -__all__ = ["Block", "MaskedSoftmax", "RollingHead", "ScoringHead", "ValueHead"] +__all__ = ["Block", "MaskedSoftmax", "ResidualBlock", "RollingHead", "ScoringHead", "ValueHead"] diff --git a/src/yahtzee_agent/modules/residual_block.py b/src/yahtzee_agent/modules/residual_block.py new file mode 100644 index 0000000..6dbbeaf --- /dev/null +++ b/src/yahtzee_agent/modules/residual_block.py @@ -0,0 +1,53 @@ +import torch +from torch import nn + + +class ResidualBlock(nn.Module): + """A residual MLP block with skip connection added before normalization. + + Architecture: x -> Linear -> Activation -> (x + residual) -> LayerNorm -> Dropout + + This design helps prevent gradient decoherence in shared representations + by providing a direct path for gradients to flow through the network. + """ + + def __init__( + self, + in_features: int, + out_features: int, + dropout_rate: float, + activation: type[nn.Module] = nn.PReLU, + ): + super().__init__() + self.linear = nn.Linear(in_features, out_features) + self.activation = activation() + self.norm = nn.LayerNorm(out_features) + self.dropout = nn.Dropout(dropout_rate) if dropout_rate > 0.0 else None + + # Projection for residual connection if dimensions don't match + self.projection = ( + nn.Linear(in_features, out_features) if in_features != out_features else None + ) + + def __call__(self, x: torch.Tensor) -> torch.Tensor: + """Call method to enable direct calls to the block.""" + return self.forward(x) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Forward pass through the residual block.""" + # Compute residual path + residual: torch.Tensor = x if self.projection is None else self.projection(x) + + # Main path: Linear -> Activation + out: torch.Tensor = self.linear(x) + out = self.activation(out) + + # Add residual connection before normalization + out = out + residual + + # Normalize and optionally dropout + out = self.norm(out) + if self.dropout is not None: + out = self.dropout(out) + + return out