Fix sb3_contrib/ars/policies.py type hint

[qgallouedec]

Description

Requires DLR-RM/stable-baselines3#1188 to be merged

Context

• I have raised an issue to propose this change (required)

Types of changes

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

Checklist:

• I've read the CONTRIBUTION guide (required)
• The functionality/performance matches that of the source (required for new training algorithms or training-related features).
• I have updated the tests accordingly (required for a bug fix or a new feature).
• I have included an example of using the feature (required for new features).
• I have included baseline results (required for new training algorithms or training-related features).
• I have updated the documentation accordingly.
• I have updated the changelog accordingly (required).
• I have reformatted the code using make format (required)
• I have checked the codestyle using make check-codestyle and make lint (required)
• I have ensured make pytest and make type both pass. (required)

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

[araffin]

Requires DLR-RM/stable-baselines3#1188 to be merged

and to be released too

[araffin]

One thing that need to be checked before merging: does it break pre-trained models? (we might have to do some manual renaming in case it does)

[qgallouedec]

All rl-baselines3-zoo/tests/test_trained_agents passed, except CartPole-v1. I'm investigating.

[qgallouedec]

Probably because ARSLinearPolicy.action_net is turned from

Linear(in_features=4, out_features=2, bias=False)

to

Sequential(
  (0): Linear(in_features=4, out_features=2, bias=False)
)

[qgallouedec]

from sb3_contrib import ARS
from torch import nn

FOLDER = "rl-trained-agents"
env_id = "CartPole-v1"

for model_name in ["best_model", env_id]:
    model = ARS.load(f"{FOLDER}/ars/{env_id}_1/{model_name}.zip")
    model.policy.action_net = nn.Sequential(model.policy.action_net)
    model.save(f"{FOLDER}/ars/{env_id}_1/{model_name}.zip")

It seems to have worked. Does that look right to you?

[qgallouedec]

For HF, I think this should also work, but I'm not familiar enough with HF Hub to be sure:

from huggingface_sb3 import load_from_hub, push_to_hub
from torch import nn
from sb3_contrib import ARS

checkpoint = load_from_hub(
    repo_id="sb3/ars-CartPole-v1",
    filename="ars-CartPole-v1.zip",
)
model = ARS.load(checkpoint)
model.policy.action_net = nn.Sequential(model.policy.action_net)

push_to_hub(
    repo_id="sb3/ars-CartPole-v1",
    filename="ars-CartPole-v1.zip",
    commit_message="Update action_net structure",
)

[araffin]

All rl-baselines3-zoo/tests/test_trained_agents passed, except CartPole-v1. I'm investigating.

I fixed that in my latest commit ;) (the name of the commit is wrong though ^^ read "load" instead of "save")

[araffin]

LGTM, please check what I did before merging ;)

[qgallouedec]

Looks good, I just turned the docstring of ARS.set_parameters into a comment, so that the true docstring is inherited from BaseAlgorithm.