[Bug]: Cannot load pre-trained policy

[araffin]

🐛 Bug

Bug introduced in #1148

We might also want to check we are not saving the weights of the features extractor two times.

To Reproduce

Just load a pretrained PPO agent saved with SB3 v1.6.2 with a CNN policy.

python enjoy.py --algo a2c --env BreakoutNoFrameskip-v4

Found by updating the RL Zoo: https://github.com/DLR-RM/rl-baselines3-zoo/actions/runs/3760560053/jobs/6391404574

Relevant log output / Error message

RuntimeError: Error(s) in loading state_dict for ActorCriticCnnPolicy:
	Missing key(s) in state_dict: "pi_features_extractor.cnn.0.weight", "pi_features_extractor.cnn.0.bias", "pi_features_extractor.cnn.2.weight", "pi_features_extractor.cnn.2.bias", "pi_features_extractor.cnn.4.weight", "pi_features_extractor.cnn.4.bias", "pi_features_extractor.linear.0.weight", "pi_features_extractor.linear.0.bias", "vf_features_extractor.cnn.0.weight", "vf_features_extractor.cnn.0.bias", "vf_features_extractor.cnn.2.weight", "vf_features_extractor.cnn.2.bias", "vf_features_extractor.cnn.4.weight", "vf_features_extractor.cnn.4.bias", "vf_features_extractor.linear.0.weight", "vf_features_extractor.linear.0.bias".

System Info

No response

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.

[araffin]

@qgallouedec one easy fix is to tell user to use exact_match=False and then save the policy again.
I'm working on a better one (or at least one that doesn't require user action).