Hotfix - Return the new updated key in function _train (#46)

* return the new updated key in _train * Add regression test and update version --------- Co-authored-by: Antonin Raffin <antonin.raffin@ensta.org>
araffin · Apr 12, 2024 · 42caa65 · 42caa65
1 parent fcd647e
commit 42caa65
Show file tree

Hide file tree

Showing 6 changed files with 8 additions and 5 deletions.
diff --git a/sbx/crossq/crossq.py b/sbx/crossq/crossq.py
@@ -474,6 +474,6 @@ def one_update(i: int, carry: Dict[str, Any]) -> Dict[str, Any]:
             update_carry["qf_state"],
             update_carry["actor_state"],
             update_carry["ent_coef_state"],
-            key,
+            update_carry["key"],
             (update_carry["info"]["actor_loss"], update_carry["info"]["qf_loss"], update_carry["info"]["ent_coef_loss"]),
         )
diff --git a/sbx/sac/sac.py b/sbx/sac/sac.py
@@ -445,6 +445,6 @@ def one_update(i: int, carry: Dict[str, Any]) -> Dict[str, Any]:
             update_carry["qf_state"],
             update_carry["actor_state"],
             update_carry["ent_coef_state"],
-            key,
+            update_carry["key"],
             (update_carry["info"]["actor_loss"], update_carry["info"]["qf_loss"], update_carry["info"]["ent_coef_loss"]),
         )
diff --git a/sbx/td3/td3.py b/sbx/td3/td3.py
@@ -335,6 +335,6 @@ def one_update(i: int, carry: Dict[str, Any]) -> Dict[str, Any]:
         return (
             update_carry["qf_state"],
             update_carry["actor_state"],
-            key,
+            update_carry["key"],
             (update_carry["info"]["actor_loss"], update_carry["info"]["qf_loss"]),
         )
diff --git a/sbx/tqc/tqc.py b/sbx/tqc/tqc.py
@@ -531,7 +531,7 @@ def one_update(i: int, carry: Dict[str, Any]) -> Dict[str, Any]:
             update_carry["qf2_state"],
             update_carry["actor_state"],
             update_carry["ent_coef_state"],
-            key,
+            update_carry["key"],
             (
                 update_carry["info"]["qf1_loss"],
                 update_carry["info"]["qf2_loss"],

diff --git a/sbx/version.txt b/sbx/version.txt
@@ -1 +1 @@
-0.14.0
+0.15.0
diff --git a/tests/test_run.py b/tests/test_run.py
@@ -84,7 +84,10 @@ def test_sac_td3(tmp_path, model_class) -> None:
         gradient_steps=1,
         learning_rate=1e-3,
     )
+    key_before_learn = model.key
     model.learn(110)
+    # See issue #45
+    assert not np.allclose(key_before_learn, model.key)
     check_save_load(model, model_class, tmp_path)