designing-feature-extractor

name	Designing Feature Extractor
description	This skill should be used when the user asks to "design a feature extractor", "build a CNN feature extractor", "implement a Nature CNN", "add a ResNet backbone to RL", "handle dictionary observations", "process image observations in RL", "build preprocessing layers", or "implement multi-modal observation handling". Do NOT use Dropout or BatchNorm inside standard RL pipelines.
version	0.1.0

Designing Feature Extractor

The feature extractor transforms raw observations (pixels, vectors, or dicts) into a latent embedding vector consumed by the Actor and Critic heads. Correct architecture selection prevents destabilization and training collapse.

Vision-Based Extraction (CNNs)

Nature-CNN (Low-Resolution / Atari-Scale)

Use for 84×84 grayscale or small RGB inputs:

import torch.nn as nn

class NatureCNN(nn.Module):
    def __init__(self, obs_shape, features_dim=512):
        super().__init__()
        n_input_channels = obs_shape[0]
        self.cnn = nn.Sequential(
            nn.Conv2d(n_input_channels, 32, kernel_size=8, stride=4, padding=0),
            nn.ReLU(),
            nn.Conv2d(32, 64, kernel_size=4, stride=2, padding=0),
            nn.ReLU(),
            nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=0),
            nn.ReLU(),
            nn.Flatten(),
        )
        # Compute flattened size dynamically
        with torch.no_grad():
            sample = torch.zeros(1, *obs_shape)
            n_flatten = self.cnn(sample).shape[1]
        self.linear = nn.Sequential(
            nn.Linear(n_flatten, features_dim),
            nn.ReLU()
        )

    def forward(self, obs):
        return self.linear(self.cnn(obs))

ResNet Backbone (High-Fidelity Inputs)

Use ResNet-18 or ResNet-34 for high-resolution or complex visual environments. Avoid ResNet-50+ unless using pre-trained weights — RL sample efficiency drops brutally with massive over-parameterization.

Strip the classification head and project to a dense flattened vector:

import torchvision.models as models

class ResNetExtractor(nn.Module):
    def __init__(self, features_dim=512):
        super().__init__()
        backbone = models.resnet18(pretrained=False)
        # Remove avgpool and fc (classification head)
        self.backbone = nn.Sequential(*list(backbone.children())[:-2])
        self.pool = nn.AdaptiveAvgPool2d((1, 1))
        self.projection = nn.Linear(512, features_dim)

    def forward(self, obs):
        x = self.backbone(obs)
        x = self.pool(x).flatten(1)
        return self.projection(x)

See references/resnet_backbone_patterns.md for detailed ResNet-18/34 architecture variants.

Vector-Based Extraction (MLPs)

For proprioceptive states (joint angles, velocities, coordinates):

class MLPExtractor(nn.Module):
    def __init__(self, obs_dim, features_dim=256):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(obs_dim, 256),
            nn.Mish(),
            nn.Linear(256, 256),
            nn.Mish(),
            nn.Linear(256, features_dim),
            nn.Mish(),
        )

    def forward(self, obs):
        return self.net(obs)

Normalization rule: Do NOT use Dropout or BatchNorm inside standard RL pipelines. They break the required deterministic property of value targets and severely destabilize off-policy learning. Use LayerNorm only if normalization is absolutely necessary:

nn.LayerNorm(features_dim)  # OK if needed
# nn.BatchNorm1d(features_dim)  # NEVER — breaks off-policy stability
# nn.Dropout(0.1)               # NEVER — breaks deterministic value targets

Multi-Modal Dictionary Observation Handling

Build dual extractors when the observation is a Dictionary (e.g., Image + Robotic Joint angles). Use the script for the standard projection builder:

• scripts/multimodal_extractor.py — Builds CNN + MLP branches and concatenates outputs

Architecture:

• CNN for Image → 256-dimensional vector
• MLP for Joints → 64-dimensional vector
• Concatenate → 320-dimensional combined feature vector

class MultiModalExtractor(nn.Module):
    def __init__(self, img_shape, joints_dim):
        super().__init__()
        self.cnn = NatureCNN(img_shape, features_dim=256)
        self.mlp = MLPExtractor(joints_dim, features_dim=64)

    def forward(self, obs_dict):
        img_features = self.cnn(obs_dict["image"])
        joint_features = self.mlp(obs_dict["joints"])
        return torch.cat([img_features, joint_features], dim=-1)  # → 320-dim

Weight Sharing Decision

Determine whether the Actor and Critic share the feature extractor weights:

Scenario	Recommendation
Simple discrete tasks (Atari DQN)	Share — reduces parameters, speeds training
Complex continuous control (PPO)	Separate — Actor gradients can destroy Critic representations
Off-policy (SAC, TD3)	Always separate — Critic is updated far more frequently

Additional Resources

Scripts

• scripts/multimodal_extractor.py — Multi-modal dictionary projection builder with CNN + MLP branches

References

• references/resnet_backbone_patterns.md — ResNet-18/34 backbone configurations, pre-training strategies