Merge branch 'master' into fix/car-racing-v3

araffin · web-flow · commit c6d3b0eda693 · 2025-06-16T11:13:48.000+02:00
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
@@ -53,6 +53,7 @@ jobs:
         run: |
           uv pip install --system gymnasium==${{ matrix.gymnasium-version }}
           uv pip install --system "numpy<2"
+          uv pip install --system "ale-py==0.10.1"
         # Only run for python 3.10, downgrade gym to 0.29.1
         if: matrix.gymnasium-version != '1.0.0'
 
diff --git a/.github/workflows/trained_agents.yml b/.github/workflows/trained_agents.yml
@@ -49,14 +49,12 @@ jobs:
           # Use headless version
           uv pip install --system opencv-python-headless
           uv pip install --system -e .[plots,tests]
-          # Downgrade numpy to run pybullet agents
-          # See https://github.com/bulletphysics/bullet3/issues/4649
-          uv pip install --system "numpy<2"
 
       - name: Install specific version of gym
         run: |
           uv pip install --system gymnasium==${{ matrix.gymnasium-version }}
           uv pip install --system "numpy<2"
+          uv pip install --system "ale-py==0.10.1"
         # Only run for python 3.10, downgrade gym to 0.29.1
         if: matrix.gymnasium-version != '1.0.0'
 
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -1,3 +1,20 @@
+## Release 2.6.1 (WIP)
+
+### Breaking Changes
+- Upgraded to SB3 >= 2.6.1
+- `linear_schedule` now returns a `SimpleLinearSchedule` object for better portability
+- Renamed `LunarLander-v2` to `LunarLander-v3` in hyperparameters
+
+### New Features
+
+### Bug fixes
+- Docker GPU images are now working again
+- Use `ConstantSchedule`, and `SimpleLinearSchedule` instead of `constant_fn` and `linear_schedule`
+
+### Documentation
+
+### Other
+
 ## Release 2.6.0 (2025-03-24)
 
 ### Breaking Changes
diff --git a/docker/Dockerfile b/docker/Dockerfile
@@ -1,4 +1,4 @@
-ARG PARENT_IMAGE
+ARG PARENT_IMAGE=stablebaselines/stable-baselines3
 FROM $PARENT_IMAGE
 
 USER root
diff --git a/hyperparams/a2c.yml b/hyperparams/a2c.yml
@@ -19,7 +19,7 @@ CartPole-v1:
   policy: 'MlpPolicy'
   ent_coef: 0.0
 
-LunarLander-v2:
+LunarLander-v3:
   n_envs: 8
   n_timesteps: !!float 2e5
   policy: 'MlpPolicy'
diff --git a/hyperparams/ars.yml b/hyperparams/ars.yml
@@ -19,7 +19,7 @@ Pendulum-v1: &pendulum-params
   zero_policy: False
 
 # TO BE Tuned
-LunarLander-v2:
+LunarLander-v3:
   <<: *pendulum-params
   n_delta: 6
   n_top: 1
diff --git a/hyperparams/dqn.yml b/hyperparams/dqn.yml
@@ -50,7 +50,7 @@ MountainCar-v0:
   policy_kwargs: "dict(net_arch=[256, 256])"
 
 # Tuned
-LunarLander-v2:
+LunarLander-v3:
   n_timesteps: !!float 1e5
   policy: 'MlpPolicy'
   learning_rate: !!float 6.3e-4
diff --git a/hyperparams/ppo.yml b/hyperparams/ppo.yml
@@ -111,7 +111,7 @@ BipedalWalkerHardcore-v3:
   learning_rate: lin_2.5e-4
   clip_range: lin_0.2
 
-LunarLander-v2:
+LunarLander-v3:
   n_envs: 16
   n_timesteps: !!float 1e6
   policy: 'MlpPolicy'
diff --git a/hyperparams/qrdqn.yml b/hyperparams/qrdqn.yml
@@ -43,7 +43,7 @@ MountainCar-v0:
   policy_kwargs: "dict(net_arch=[256, 256], n_quantiles=25)"
 
 # Tuned
-LunarLander-v2:
+LunarLander-v3:
   n_timesteps: !!float 1e5
   policy: 'MlpPolicy'
   learning_rate: lin_1.5e-3
diff --git a/hyperparams/trpo.yml b/hyperparams/trpo.yml
@@ -23,7 +23,7 @@ Pendulum-v1:
   sde_sample_freq: 4
 
 # Tuned
-LunarLander-v2:
+LunarLander-v3:
   n_envs: 2
   n_timesteps: !!float 2e5
   policy: 'MlpPolicy'
diff --git a/requirements.txt b/requirements.txt
@@ -1,5 +1,5 @@
 gym==0.26.2
-stable-baselines3[extra,tests,docs]>=2.6.0,<3.0
+stable-baselines3[extra,tests,docs]>=2.6.1a1,<3.0
 box2d-py==2.3.8
 pybullet_envs_gymnasium>=0.6.0
 # minigrid
diff --git a/rl_zoo3/exp_manager.py b/rl_zoo3/exp_manager.py
@@ -50,7 +50,14 @@
 import rl_zoo3.import_envs  # noqa: F401
 from rl_zoo3.callbacks import SaveVecNormalizeCallback, TrialEvalCallback
 from rl_zoo3.hyperparams_opt import HYPERPARAMS_CONVERTER, HYPERPARAMS_SAMPLER
-from rl_zoo3.utils import ALGOS, get_callback_list, get_class_by_name, get_latest_run_id, get_wrapper_class, linear_schedule
+from rl_zoo3.utils import (
+    ALGOS,
+    SimpleLinearSchedule,
+    get_callback_list,
+    get_class_by_name,
+    get_latest_run_id,
+    get_wrapper_class,
+)
 
 
 class ExperimentManager:
@@ -381,7 +388,7 @@ def _preprocess_schedules(hyperparams: dict[str, Any]) -> dict[str, Any]:
             if isinstance(hyperparams[key], str):
                 schedule, initial_value = hyperparams[key].split("_")
                 initial_value = float(initial_value)
-                hyperparams[key] = linear_schedule(initial_value)
+                hyperparams[key] = SimpleLinearSchedule(initial_value)
             elif isinstance(hyperparams[key], (float, int)):
                 # Negative value: ignore (ex: for clipping)
                 if hyperparams[key] < 0:
diff --git a/rl_zoo3/utils.py b/rl_zoo3/utils.py
@@ -292,25 +292,33 @@ def make_env(**kwargs) -> gym.Env:
     return env
 
 
-def linear_schedule(initial_value: Union[float, str]) -> Callable[[float], float]:
+class SimpleLinearSchedule:
+    """
+    Linear learning rate schedule (from initial value to zero),
+    simpler than sb3 LinearSchedule.
+
+    :param initial_value: (float or str) The initial value for the schedule
+    """
+
+    def __init__(self, initial_value: Union[float, str]) -> None:
+        # Force conversion to float
+        self.initial_value = float(initial_value)
+
+    def __call__(self, progress_remaining: float) -> float:
+        return progress_remaining * self.initial_value
+
+    def __repr__(self) -> str:
+        return f"SimpleLinearSchedule(initial_value={self.initial_value})"
+
+
+def linear_schedule(initial_value: Union[float, str]) -> SimpleLinearSchedule:
     """
     Linear learning rate schedule.
 
     :param initial_value: (float or str)
-    :return: (function)
+    :return: A `SimpleLinearSchedule` object
     """
-    # Force conversion to float
-    initial_value_ = float(initial_value)
-
-    def func(progress_remaining: float) -> float:
-        """
-        Progress will decrease from 1 (beginning) to 0
-        :param progress_remaining: (float)
-        :return: (float)
-        """
-        return progress_remaining * initial_value_
-
-    return func
+    return SimpleLinearSchedule(initial_value)
 
 
 def get_trained_models(log_folder: str) -> dict[str, tuple[str, str]]:
diff --git a/rl_zoo3/version.txt b/rl_zoo3/version.txt
@@ -1 +1 @@
-2.6.0
+2.6.1a1
diff --git a/rl_zoo3/wrappers.py b/rl_zoo3/wrappers.py
@@ -311,8 +311,8 @@ class MaskVelocityWrapper(gym.ObservationWrapper):
         "MountainCar-v0": np.array([1]),
         "MountainCarContinuous-v0": np.array([1]),
         "Pendulum-v1": np.array([2]),
-        "LunarLander-v2": np.array([2, 3, 5]),
-        "LunarLanderContinuous-v2": np.array([2, 3, 5]),
+        "LunarLander-v3": np.array([2, 3, 5]),
+        "LunarLanderContinuous-v3": np.array([2, 3, 5]),
     }
 
     def __init__(self, env: gym.Env):
diff --git a/scripts/create_mujoco_jobs.py b/scripts/create_mujoco_jobs.py
@@ -49,7 +49,7 @@
             ]
             args = list(map(str, args))
 
-            command = " ".join(["python", "-u", "train.py", *args])
+            command = " ".join(["python", "-u", "train.py", *args])  # type: ignore[list-item]
 
             ok = subprocess.call(["sbatch", "cluster_torchy.sh", algo, env_id, "ablation", command])
             time.sleep(0.05)
diff --git a/scripts/run_docker_gpu.sh b/scripts/run_docker_gpu.sh
@@ -8,6 +8,6 @@ echo $cmd_line
 
 # Note: --user=root is needed, as the current user id/group id will be used
 # to mount the log directory (and $MAMBAUSER is not root)
-docker run -it --user=root --runtime=nvidia --rm --network host --ipc=host \
+docker run -it --user=root --gpus all --rm --network host --ipc=host \
   --mount src=$(pwd),target=/home/mambauser/code/rl_zoo3,type=bind stablebaselines/rl-baselines3-zoo:latest\
   bash -c "cd /home/mambauser/code/rl_zoo3/ && $cmd_line"
diff --git a/setup.py b/setup.py
@@ -15,7 +15,7 @@
 See https://github.com/DLR-RM/rl-baselines3-zoo
 """
 install_requires = [
-    "sb3_contrib>=2.6.0,<3.0",
+    "sb3_contrib>=2.6.1a1,<3.0",
     "gymnasium>=0.29.1,<1.2.0",
     "huggingface_sb3>=3.0,<4.0",
     "tqdm",
diff --git a/tests/test_wrappers.py b/tests/test_wrappers.py
@@ -1,14 +1,16 @@
 import gymnasium as gym
+import numpy as np
 import pytest
 import stable_baselines3 as sb3
+from sb3_contrib.common.wrappers import TimeFeatureWrapper
 from stable_baselines3 import A2C
 from stable_baselines3.common.env_checker import check_env
 from stable_baselines3.common.env_util import DummyVecEnv
 
 import rl_zoo3.import_envs
 import rl_zoo3.wrappers
-from rl_zoo3.utils import get_wrapper_class
-from rl_zoo3.wrappers import ActionNoiseWrapper, DelayedRewardWrapper, HistoryWrapper, TimeFeatureWrapper
+from rl_zoo3.utils import SimpleLinearSchedule, get_wrapper_class, linear_schedule
+from rl_zoo3.wrappers import ActionNoiseWrapper, DelayedRewardWrapper, HistoryWrapper
 
 
 def test_wrappers():
@@ -55,3 +57,11 @@ def test_get_vec_env_wrapper(vec_env_wrapper):
     if wrapper_class is not None:
         env = wrapper_class(env)
     A2C("MlpPolicy", env).learn(16)
+
+
+def test_linear_schedule():
+    schedule = linear_schedule(100)
+    assert isinstance(schedule, SimpleLinearSchedule)
+    assert np.allclose(schedule(1.0), 100.0)
+    assert np.allclose(schedule(0.5), 50.0)
+    assert np.allclose(schedule(0.0), 0.0)

Original file line number	Diff line number	Diff line change
`@@ -1,4 +1,4 @@`
`1`		`-ARG PARENT_IMAGE`
	`1`	`+ARG PARENT_IMAGE=stablebaselines/stable-baselines3`
`2`	`2`	`FROM $PARENT_IMAGE`
`3`	`3`
`4`	`4`	`USER root`