return the new updated key in _train

araffin · Apr 12, 2024 · 84feff5 · 84feff5
1 parent fcd647e
commit 84feff5
Show file tree

Hide file tree

Showing 4 changed files with 4 additions and 4 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"],