Optional BatchNorm integration in NatureCNN

[Mahsarnzh]

trafficstars

Description

Added an optional BatchNorm integration to the NatureCNN architecture used in the feature extractor module of Stable-Baselines3. This enhancement introduces a use_batch_norm flag to toggle Batch Normalization after each convolutional layer. This change provides a performance and stability improvement option for image-based environments.

Motivation and Context

This change will solve the exploding gradients problem and in case it is set to False it does not change anything, however if set to True it will help converge much faster and enables us to use higher learning rates. Further than that this change allows users to optionally enable Batch Normalization in NatureCNN, which can improve training stability and convergence, especially in environments with high variance in pixel input. I initially explored alternatives (LayerNorm, GroupNorm). BatchNorm showed the best trade-off of speed and stability and convergence.

N/A N/A

• [x] I have raised an issue to propose this change (#2131 for new features and bug fixes)

Types of changes

• [ ] Bug fix (non-breaking change which fixes an issue)
• [x] New feature (non-breaking change which adds functionality)
• [ ] Breaking change (fix or feature that would cause existing functionality to change)
• [x] Documentation (update in the documentation)

Checklist

• [x] I've read the CONTRIBUTION guide (required)
• [x] I have updated the changelog accordingly (required).
• [ ] My change requires a change to the documentation.
• [x] I have updated the tests accordingly (required for a bug fix or a new feature).
• [x] I have updated the documentation accordingly.
• [x] I have opened an associated PR on the SB3-Contrib repository (if necessary)
• [ ] I have opened an associated PR on the RL-Zoo3 repository (if necessary)
• [x] I have reformatted the code using make format (required)
• [x] I have checked the codestyle using make check-codestyle and make lint (required)
• [x] I have ensured make pytest and make type both pass. (required)
• [x] I have checked that the documentation builds using make doc (required)

Note: You can run most of the checks using make commit-checks.

Note: we are using a maximum length of 127 characters per line