Update Hugging Face 🤗 Push To Hub (#379)

* Update push_to_hub.py

* Only render when needed

* Patch Atari game video recording

* Patch atari rendering

* Remove Atari patch, will be fixed by SB3 update

* Fix record video steps + update comments

* Update versions

---------

Co-authored-by: Antonin Raffin <antonin.raffin@dlr.de>
Co-authored-by: Antonin Raffin <antonin.raffin@ensta.org>

Author: simoninithomas

CHANGELOG.md

@@ -1,15 +1,18 @@
-## Release 2.0.0a9 (WIP)
+## Release 2.0.0a12 (WIP)
 
 ### Breaking Changes
 - Upgraded to gym 0.26+
 - Fixed bug in HistoryWrapper, now returns the correct obs space limits
 - Upgraded to SB3 >= 2.0.0
+- Upgraded to Huggingface-SB3 >= 2.2.5
 
 ### New Features
 - Gym 0.26+ patches to continue working with pybullet and TimeLimit wrapper
 
 ### Bug fixes
-- Renamed ``CarRacing-v1`` to ``CarRacing-v2`` in hyperparameters
+- Renamed `CarRacing-v1` to `CarRacing-v2` in hyperparameters
+- Huggingface push to hub now accepts a `--n-timesteps` argument to adjust the length of the video
+- Fixed `record_video` steps (before it was stepping in a closed env)
 
 ## Release 1.8.0 (2023-04-07)
 

requirements.txt

@@ -1,6 +1,6 @@
 gym==0.26.2
-stable-baselines3[extra_no_roms,tests,docs]>=2.0.0a9
-sb3-contrib>=2.0.0a9
+stable-baselines3[extra_no_roms,tests,docs]>=2.0.0a13
+sb3-contrib>=2.0.0a13
 box2d-py==2.3.8
 pybullet
 # minigrid
@@ -14,7 +14,7 @@ plotly
 # panda-gym~=3.0.1
 rliable>=1.0.5
 wandb
-huggingface_sb3>=2.2.1
+huggingface_sb3>=2.2.5
 seaborn
 tqdm
 rich

rl_zoo3/push_to_hub.py

@@ -20,9 +20,9 @@
 from wasabi import Printer
 
 import rl_zoo3.import_envs  # noqa: F401 pylint: disable=unused-import
-from rl_zoo3 import ALGOS, create_test_env, get_saved_hyperparams
+from rl_zoo3 import ALGOS, get_saved_hyperparams
 from rl_zoo3.exp_manager import ExperimentManager
-from rl_zoo3.utils import StoreDict, get_model_path
+from rl_zoo3.utils import StoreDict, create_test_env, get_model_path
 
 msg = Printer()
 
@@ -277,12 +277,12 @@ def package_to_hub(
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
-    parser.add_argument("--env", help="environment ID", type=EnvironmentName, required=True)
+    parser.add_argument("--env", help="Environment ID", type=EnvironmentName, required=True)
     parser.add_argument("-f", "--folder", help="Log folder", type=str, required=True)
     parser.add_argument("--algo", help="RL Algorithm", type=str, required=True, choices=list(ALGOS.keys()))
-    parser.add_argument("-n", "--n-timesteps", help="number of timesteps", default=1000, type=int)
+    parser.add_argument("-n", "--n-timesteps", help="Number of timesteps for the video recording", default=1000, type=int)
     parser.add_argument("--num-threads", help="Number of threads for PyTorch (-1 to use default)", default=-1, type=int)
-    parser.add_argument("--n-envs", help="number of environments", default=1, type=int)
+    parser.add_argument("--n-envs", help="Number of environments", default=1, type=int)
     parser.add_argument("--exp-id", help="Experiment ID (default: 0: latest, -1: no exp folder)", default=0, type=int)
     parser.add_argument("--verbose", help="Verbose mode (0: no output, 1: INFO)", default=1, type=int)
     parser.add_argument(
@@ -357,6 +357,12 @@ def package_to_hub(
             loaded_args = yaml.load(f, Loader=yaml.UnsafeLoader)  # pytype: disable=module-attr
             if loaded_args["env_kwargs"] is not None:
                 env_kwargs = loaded_args["env_kwargs"]
+
+    # render and record video by default
+    should_render = not args.no_render
+    if should_render:
+        env_kwargs.update(render_mode="rgb_array")
+
     # overwrite with command line arguments
     if args.env_kwargs is not None:
         env_kwargs.update(args.env_kwargs)
@@ -367,7 +373,7 @@ def package_to_hub(
         stats_path=maybe_stats_path,
         seed=args.seed,
         log_dir=None,
-        should_render=not args.no_render,
+        should_render=should_render,
         hyperparams=deepcopy(hyperparams),
         env_kwargs=env_kwargs,
     )
@@ -377,6 +383,12 @@ def package_to_hub(
         # Dummy buffer size as we don't need memory to enjoy the trained agent
         kwargs.update(dict(buffer_size=1))
 
+        # Hack due to breaking change in v1.6
+        # handle_timeout_termination cannot be at the same time
+        # with optimize_memory_usage
+        if "optimize_memory_usage" in hyperparams:
+            kwargs.update(optimize_memory_usage=False)
+
     # Note: we assume that we push models using the same machine (same python version)
     # that trained them, if not, we would need to pass custom object as in enjoy.py
     custom_objects: Dict[str, Any] = {}
@@ -411,6 +423,6 @@ def package_to_hub(
         n_eval_episodes=10,
         token=None,
         local_repo_path="hub",
-        video_length=1000,
+        video_length=args.n_timesteps,
         generate_video=not args.no_render,
     )

rl_zoo3/record_video.py

@@ -13,12 +13,12 @@
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
-    parser.add_argument("--env", help="environment ID", type=EnvironmentName, default="CartPole-v1")
+    parser.add_argument("--env", help="Environment ID", type=EnvironmentName, default="CartPole-v1")
     parser.add_argument("-f", "--folder", help="Log folder", type=str, default="rl-trained-agents")
     parser.add_argument("-o", "--output-folder", help="Output folder", type=str)
     parser.add_argument("--algo", help="RL Algorithm", default="ppo", type=str, required=False, choices=list(ALGOS.keys()))
-    parser.add_argument("-n", "--n-timesteps", help="number of timesteps", default=1000, type=int)
-    parser.add_argument("--n-envs", help="number of environments", default=1, type=int)
+    parser.add_argument("-n", "--n-timesteps", help="Number of timesteps", default=1000, type=int)
+    parser.add_argument("--n-envs", help="Number of environments", default=1, type=int)
     parser.add_argument("--deterministic", action="store_true", default=False, help="Use deterministic actions")
     parser.add_argument("--stochastic", action="store_true", default=False, help="Use stochastic actions")
     parser.add_argument("--seed", help="Random generator seed", type=int, default=0)
@@ -150,7 +150,7 @@
     lstm_states = None
     episode_starts = np.ones((env.num_envs,), dtype=bool)
     try:
-        for _ in range(video_length + 1):
+        for _ in range(video_length):
             action, lstm_states = model.predict(
                 obs,  # type: ignore[arg-type]
                 state=lstm_states,

rl_zoo3/version.txt

@@ -1 +1 @@
-2.0.0a9
+2.0.0a13

setup.py

@@ -27,9 +27,9 @@
     },
     entry_points={"console_scripts": ["rl_zoo3=rl_zoo3.cli:main"]},
     install_requires=[
-        "sb3_contrib>=2.0.0a9",
+        "sb3_contrib>=2.0.0a13",
         "gym==0.26.2",
-        "huggingface_sb3>=2.2.1",
+        "huggingface_sb3>=2.2.5",
         "tqdm",
         "rich",
         "optuna",