add residuals in the trunk

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)
 

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

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