Fix saving jitted model and simplify load

[araffin]

Description

closes #2137

Motivation and Context

• I have raised an issue to propose this change (required for new features and bug fixes)

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)
• I have updated the changelog accordingly (required).
• My change requires a change to the documentation.
• I have updated the tests accordingly (required for a bug fix or a new feature).
• I have updated the documentation accordingly.
• I have opened an associated PR on the SB3-Contrib repository (if necessary)
• I have opened an associated PR on the RL-Zoo3 repository (if necessary)
• 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)
• I have checked that the documentation builds using make doc (required)

Note: You can run most of the checks using make commit-checks.

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

[Copilot]

Pull request overview

This PR fixes issues with saving and loading jitted (compiled) models and streamlines the loading process by utilizing PyTorch's built-in capabilities. The changes address issue #2137 by properly handling torch.compile wrapped models and simplifying the zip file loading logic.

Key changes:

• Modified get_parameters() to extract state dicts from the original model when dealing with compiled models
• Simplified load_from_zip_file() to use PyTorch's direct loading with weights_only=True
• Added test coverage for saving/loading compiled models and reduced training timesteps across tests

Reviewed changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated 1 comment.

File	Description
stable_baselines3/common/base_class.py	Updated get_parameters() to handle compiled models by extracting state dict from _orig_mod
stable_baselines3/common/save_util.py	Simplified zip file loading by removing BytesIO workaround and enabling weights_only=True
tests/test_save_load.py	Added test for compiled model save/load, reduced timesteps, and added PPO-specific parameters

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

The test verifies that loading works after saving a compiled model, but doesn't verify the loaded model's behavior. Consider asserting that the loaded model produces the same predictions or has the same state dict as before saving.