如何在带自定义特征提取器的Stable-Baselines3中替换MlpExtractor激活函数

问题背景 #

已实现自定义特征提取器NatureCNN并接入PPO的CnnPolicy，需求是仅替换MlpExtractor内的激活函数为自定义CustomActivation，同时保留原网络层布局（如默认的64-64线性层结构），尝试官方文档方法后报错，需可行实现方案。

解决方案 #

核心思路是自定义继承自原MlpExtractor的子类，重写网络构建逻辑，仅替换激活函数，其余结构完全复用默认设置，再通过policy_kwargs传入该自定义类。

步骤1：完善自定义激活函数 #

确保CustomActivation可正常实例化并运行：

import torch.nn as nn
import torch as th

class CustomActivation(nn.Module):
    def __init__(self, param=None):
        super(CustomActivation, self).__init__()
        # 初始化自定义激活的参数（根据你的需求调整）
        self.param = param if param is not None else 1.0

    def forward(self, x):
        # 实现自定义激活逻辑，示例用带参数的sigmoid变体
        return th.sigmoid(x * self.param)

步骤2：自定义MlpExtractor子类 #

继承原MlpExtractor，重写网络创建逻辑，替换激活函数：

from stable_baselines3.common.torch_layers import MlpExtractor

class CustomMlpExtractor(MlpExtractor):
    def __init__(
        self,
        feature_dim: int,
        net_arch: dict,
        activation_fn: nn.Module,
        device: th.device
    ):
        # 调用父类初始化，保留原有参数逻辑
        super().__init__(feature_dim, net_arch, activation_fn, device)
        
        # 重新构建policy_net和value_net，替换激活函数为CustomActivation
        # 对应你期望的结构：Linear(512,64) -> CustomActivation -> Linear(64,64) -> CustomActivation
        self.policy_net = nn.Sequential(
            nn.Linear(feature_dim, 64),
            CustomActivation(),  # 替换原Tanh
            nn.Linear(64, 64),
            CustomActivation()   # 替换原Tanh
        ).to(device)
        
        self.value_net = nn.Sequential(
            nn.Linear(feature_dim, 64),
            CustomActivation(),  # 替换原Tanh
            nn.Linear(64, 64),
            CustomActivation()   # 替换原Tanh
        ).to(device)
        
        # shared_net保持为空（和默认结构一致）
        self.shared_net = nn.Sequential()

步骤3：修改policy_kwargs并初始化模型 #

在原有policy_kwargs基础上，添加自定义mlp_extractor_class指定我们的CustomMlpExtractor：

from stable_baselines3 import PPO
from your_module import NatureCNN  # 替换为你的NatureCNN所在模块

policy_kwargs = dict(
    features_extractor_class=NatureCNN,
    mlp_extractor_class=CustomMlpExtractor,
    # 如果需要给CustomActivation传参，可通过activation_fn传递实例
    # activation_fn=lambda: CustomActivation(param=0.5)
)

model = PPO(
    'CnnPolicy', 
    env, 
    batch_size=128,
    clip_range=0.10,
    max_grad_norm=0.5,
    verbose=1, 
    seed=1,
    device="cuda",
    tensorboard_log="./tb_logs/",
    policy_kwargs=policy_kwargs,
)

验证结构 #

初始化模型后，可打印模型结构确认：

print(model.policy)

输出的mlp_extractor部分应与你期望的结构一致，仅激活函数替换为CustomActivation。

关键注意事项 #

• 自定义MlpExtractor时，严格对齐默认结构的层维度（如64-64线性层），避免维度不匹配报错
• 如果CustomActivation需要参数，可通过policy_kwargs中的activation_fn传递实例（如activation_fn=lambda: CustomActivation(param=2.0)）
• 确保所有模块都移动到指定设备（如cuda），避免张量设备不匹配问题

内容的提问来源于stack exchange，提问作者Vatsal Aggarwal