python310Packages.bsuite: Enable more tests (#181827)

Co-authored-by: Sandro <sandro.jaeckel@gmail.com>

pkgs/development/python-modules/bsuite/default.nix

@@ -20,9 +20,12 @@
 , trfl
 , optax
 , pytestCheckHook
-, dm-sonnet }:
+, dm-sonnet
+, rlax
+, distrax
+}:
 
-buildPythonPackage rec {
+let bsuite = buildPythonPackage rec {
   pname = "bsuite";
   version = "0.3.5";
 
@@ -49,10 +52,12 @@ buildPythonPackage rec {
   ];
 
   checkInputs = [
+    distrax
     dm-haiku
     dm-sonnet
     optax
     pytestCheckHook
+    rlax
     tensorflow-probability
     trfl
   ];
@@ -61,15 +66,6 @@ buildPythonPackage rec {
     "bsuite"
   ];
 
-  disabledTestPaths = [
-    # Disabled because tests require module rlax but this results in infinite
-    # recursion error
-    "bsuite/baselines/jax/actor_critic/run_test.py"
-    "bsuite/baselines/jax/actor_critic_rnn/run_test.py"
-    "bsuite/baselines/jax/boot_dqn/run_test.py"
-    "bsuite/baselines/jax/dqn/run_test.py"
-  ];
-
   disabledTests = [
     # Tests require network connection
     "test_run9"
@@ -89,6 +85,13 @@ buildPythonPackage rec {
     "test_episode_truncation"
   ];
 
+  # escape infinite recursion with rlax
+  doCheck = false;
+
+  passthru.tests = {
+    check = bsuite.overridePythonAttrs (_: { doCheck = true; });
+  };
+
   meta = with lib; {
     description = ''
       Core RL Behaviour Suite. A collection of reinforcement learning
@@ -98,4 +101,4 @@ buildPythonPackage rec {
     license = licenses.asl20;
     maintainers = with maintainers; [ onny ];
   };
-}
+}; in bsuite