[Bug] The entropy coefficient remains constant throughout training when the model is loaded from a saved checkpoint model

[Atlis]

🐛 Bug

The entropy coefficient (ent_coef) remains constant throughout training when the model is loaded from a saved checkpoint model.
This is not expected since ent_coef was set to "auto".

To Reproduce

Steps to reproduce the behavior:

1. Train a model with the TQC algo with the ent_coef set to "auto".
2. Observe that the ent_coef value varies (in the Tensorboard interface).
3. Save the trained model to a file using the CheckpointCallback.
4. Stop the training.
5. Start a new training using a saved checkpoint making sure ent_coef is still set to "auto".
6. Observe that the ent_coef value remain constant throughout training (in the Tensorboard interface).

Checklist

• I have checked that there is no similar issue in the repo (required)
• I have read the documentation (required)
• I have provided a minimal working example to reproduce the bug (required) (described in the steps to reproduce above)

[araffin]

Hello,

thanks for reporting the issue =)

I could reproduce the bug using the following code:

from sb3_contrib import TQC

policy_kwargs = dict(net_arch=[64], n_critics=1)
model = TQC("MlpPolicy", "Pendulum-v0", policy_kwargs=policy_kwargs, verbose=1, seed=1)
model.learn(2000)
ent_coef = model.log_ent_coef
env = model.get_env()
model.save("tqc_pendulum")

del model

model = TQC.load("tqc_pendulum", env=env)
model.learn(2000)
new_ent_coef = model.log_ent_coef

print(ent_coef, new_ent_coef)

the issue seems to be also present for SAC (as TQC is mostly SAC + distributional RL).
I will take a look soon.

next time, please provide a minimal working example and your config (steps to reproduce is ok-ish but a minimal code example is easier for us to reproduce and debug ;) ).

[Atlis]

Ok, I take note of this. Thank you very much!

[araffin]

The issue seems to come from the ent coef optimizer which is not linked to the correct variable after loading.

[araffin]

I pushed a fix here: #392