[bug-fix] Fix save/restore critic, add test

[ervteng]

Proposed change(s)

Add critic to the list of modules for optimizer, so that it is saved/restored properly. This was introduced in #4939 and doesn't affect older releases of ML-Agents.

Types of change(s)

• Bug fix
• New feature
• Code refactor
• Breaking change
• Documentation update
• Other (please describe)

Checklist

• Added tests that prove my fix is effective or that my feature works
• Updated the changelog (if applicable)
• Updated the documentation (if applicable)
• Updated the migration guide (if applicable)

Other comments

[dongruoping]

Is this able to run on GPU? In _compare_two_policies we move the actors explicitly and I just saw that we also added self.actor.to(default_device()) to TorchPolicy. I didn't see that in optimizers and I wonder if we need to do the same thing as well.

[dongruoping]

is it intentional that optimizer2 is using policy not policy2?

[ervteng]

No this be policy2 - good catch

[ervteng]

Is this able to run on GPU? In _compare_two_policies we move the actors explicitly and I just saw that we also added self.actor.to(default_device()) to TorchPolicy. I didn't see that in optimizers and I wonder if we need to do the same thing as well.

I think I need to add this too, but the different optimizers have different components. I'll try iterating through the get_modules() and moving each of those to the default device.

edit Never mind, this should work b/c the optimizer moves itself to the default device on initialization. So if it fails, it means the optimizer itself is broken.