Adding dropout to a custom network

[Karlheinzniebuhr]

❓ Question

I'm trying to add dropout to my network, but everything I tried gets me some kind of error.

import gym
import torch
import torch.nn as nn
from stable_baselines3.common.policies import ActorCriticPolicy
from stable_baselines3 import PPO

# Define a minimal Gym environment (CartPole-v1)
envs = gym.make("CartPole-v1")

class CustomMlpPolicy(ActorCriticPolicy):
    def __init__(self, *args, **kwargs):
        super(CustomMlpPolicy, self).__init__(*args, **kwargs)
        
        # Modify the policy network to add dropout layers
        self.policy_net = nn.Sequential(
            nn.Linear(self.mlp_extractor.output_dim, 128),
            nn.Tanh(),
            nn.Dropout(p=0.5),  # Add dropout to the first hidden layer
            nn.Linear(128, 128),
            nn.Tanh(),
        )

        # Value network remains the same
        self.value_net = nn.Sequential(
            nn.Linear(self.mlp_extractor.output_dim, 256),
            nn.Tanh(),
            nn.Linear(256, 256),
            nn.Tanh(),
            nn.Linear(256, 256),
            nn.Tanh(),
            nn.Linear(256, 256),
            nn.Tanh(),
        )

# Create the PPO model using your custom policy
model = PPO(CustomMlpPolicy, envs)
model.learn(total_timesteps=10000)  # Train for a specified number of timesteps

AttributeError: 'MlpExtractor' object has no attribute 'output_dim'

I checked the documentation but it does not contain any reference to dropouts.

How can I implement dropouts?

Checklist

• I have checked that there is no similar issue in the repo
• I have read the documentation
• If code there is, it is minimal and working
• If code there is, it is formatted using the markdown code blocks for both code and stack traces.

[araffin]

Duplicate of #1069 (comment), a working example can be found in #1036

[Karlheinzniebuhr]

I'm afraid #1036 only works for SAC() since PPO didn't get an update with that PR.

[araffin]

I'm afraid #1036 only works for SAC() since PPO didn't get an update with that PR.

well, nothing prevent you from forking and doing the update yourself ;)
i mostly wanted to link a working example with dropout.