From 264a34a3cc8f85999b4886efca61649ef567704c Mon Sep 17 00:00:00 2001
From: Gaetan Lepage <gaetan@glepage.com>
Date: Sat, 28 Jun 2025 16:32:18 +0200
Subject: [PATCH 1/2] python3Packages.gymnasium: 1.1.1 -> 1.2.0

Diff: https://github.com/Farama-Foundation/gymnasium/compare/refs/tags/v1.1.1...refs/tags/v1.2.0

Changelog: https://github.com/Farama-Foundation/Gymnasium/releases/tag/v1.2.0
---
 .../python-modules/gymnasium/default.nix        | 17 ++++++-----------
 1 file changed, 6 insertions(+), 11 deletions(-)

diff --git a/pkgs/development/python-modules/gymnasium/default.nix b/pkgs/development/python-modules/gymnasium/default.nix
index e2a18f356ce8..d9088c43d27c 100644
--- a/pkgs/development/python-modules/gymnasium/default.nix
+++ b/pkgs/development/python-modules/gymnasium/default.nix
@@ -16,6 +16,7 @@
   importlib-metadata,
 
   # tests
+  array-api-compat,
   dill,
   flax,
   jax,
@@ -28,11 +29,12 @@
   pygame,
   pytestCheckHook,
   scipy,
+  torch,
 }:
 
 buildPythonPackage rec {
   pname = "gymnasium";
-  version = "1.1.1";
+  version = "1.2.0";
 
   pyproject = true;
 
@@ -40,7 +42,7 @@ buildPythonPackage rec {
     owner = "Farama-Foundation";
     repo = "gymnasium";
     tag = "v${version}";
-    hash = "sha256-5uE6ANOxVCeV5GMDGG+0j5JY2t++jw+mZFFHGl+sTfw=";
+    hash = "sha256-fQsz1Qpef9js+iqkqbfxrTQgcZT+JKjwpEiWewju2Dc=";
   };
 
   build-system = [ setuptools ];
@@ -55,6 +57,7 @@ buildPythonPackage rec {
   pythonImportsCheck = [ "gymnasium" ];
 
   nativeCheckInputs = [
+    array-api-compat
     dill
     flax
     jax
@@ -67,6 +70,7 @@ buildPythonPackage rec {
     pygame
     pytestCheckHook
     scipy
+    torch
   ];
 
   # if `doCheck = true` on Darwin, `jaxlib` is evaluated, which is both
@@ -94,15 +98,6 @@ buildPythonPackage rec {
   '';
 
   disabledTests = [
-    # Fails since jax 0.6.0
-    # Fixed on master https://github.com/Farama-Foundation/Gymnasium/commit/94019feee1a0f945b9569cddf62780f4e1a224a5
-    # TODO: un-skip at the next release
-    "test_all_env_api"
-    "test_env_determinism_rollout"
-    "test_jax_to_numpy_wrapper"
-    "test_pickle_env"
-    "test_roundtripping"
-
     # Succeeds for most environments but `test_render_modes[Reacher-v4]` fails because it requires
     # OpenGL access which is not possible inside the sandbox.
     "test_render_mode"

From c4a6eaa3241fe0ed34605cfff31722893c99df69 Mon Sep 17 00:00:00 2001
From: Gaetan Lepage <gaetan@glepage.com>
Date: Sat, 28 Jun 2025 17:07:59 +0200
Subject: [PATCH 2/2] python3Packages.stable-baselines3: patch failing test

---
 .../python-modules/stable-baselines3/default.nix           | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/pkgs/development/python-modules/stable-baselines3/default.nix b/pkgs/development/python-modules/stable-baselines3/default.nix
index 37ea041b1c79..2a5e9091ff37 100644
--- a/pkgs/development/python-modules/stable-baselines3/default.nix
+++ b/pkgs/development/python-modules/stable-baselines3/default.nix
@@ -32,6 +32,13 @@ buildPythonPackage rec {
     hash = "sha256-VnoQ8cKqPcZPpR9c3M6xJDdG7gnO9fxIa4v2kxd9Nzg=";
   };
 
+  postPatch =
+    # Environment version v0 for `CliffWalking` is deprecated
+    ''
+      substituteInPlace "tests/test_vec_normalize.py" \
+        --replace-fail "CliffWalking-v0" "CliffWalking-v1"
+    '';
+
   build-system = [ setuptools ];
 
   pythonRelaxDeps = [