DLR-RM/stable-baselines3

[Enhancement]: Wrong gains for weight initialization

Open

#1,559 opened on Jun 16, 2023

View on GitHub
 (2 comments) (0 reactions) (1 assignee)Python (6,550 stars) (1,407 forks)batch import
enhancementhelp wanted

Description

Enhancement

The recommended gains for the weight init depend on the used activation function, see torch docs. However, as for now the used gains are statically implemented and always the same in ActorCriticPolicies. See here.

I recommend making the gains dependent on the activation function used(, i.e. probably mainly ReLU and tanh).

If you agree with this, I would like to implement it myself and PR.

Thanks and a good day!

To Reproduce

--

Relevant log output / Error message

--

System Info

--

Checklist

  • I have checked that there is no similar issue in the repo
  • I have read the documentation
  • I have provided a minimal working example to reproduce the bug
  • I've used the markdown code blocks for both code and stack traces.

Contributor guide

Tech stack
pythonpytorch
Domain
machine learning
Issue type
feature
DifficultyEstimated implementation difficulty for a new contributor, from 1 for very small changes to 5 for expert-level work.
3
Estimated timeA rough time range for an experienced contributor to investigate, implement, test, and prepare a pull request.
1-3 hours
Activity statusHow available the issue appears right now: fresh, active, stale, blocked, or waiting on maintainer input.
needs maintainer response
ClarityHow clearly the issue explains the expected change, acceptance criteria, and next step.
clear
Prerequisites
Basic knowledge of PyTorchUnderstanding of weight initializationFamiliarity with stable baselines3 codebase
Newbie friendlinessA 1-100 score estimating how approachable this issue is for first-time contributors.
60
Research direction
The issue is in stable baselines3/common/policies.py around line 589 where gains are statically set. The goal is to dynamically set gains based on activation function, following PyTorch's recommendations. The contributor should look at how activation functions are selected in the policy network and adjust the initialization accordingly. No linked PRs or assignee comments are visible.