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

   1 {
   2   lib,
   3   absl-py,
   4   buildPythonPackage,
   5   distrax,
   6   dm-env,
   7   dm-haiku,
   8   dm-sonnet,
   9   dm-tree,
  10   fetchpatch,
  11   fetchPypi,
  12   frozendict,
  13   gym,
  14   matplotlib,
  15   mizani,
  16   optax,
  17   pandas,
  18   patsy,
  19   plotnine,
  20   pytestCheckHook,
  21   pythonOlder,
  22   rlax,
  23   scikit-image,
  24   scipy,
  25   setuptools,
  26   statsmodels,
  27   tensorflow-probability,
  28   termcolor,
  29 }:
  30
  31 let
  32   bsuite = buildPythonPackage rec {
  33     pname = "bsuite";
  34     version = "0.3.5";
  35     pyproject = true;
  36
  37     disabled = pythonOlder "3.7";
  38
  39     src = fetchPypi {
  40       inherit pname version;
  41       hash = "sha256-ak9McvXl7Nz5toUaPaRaJek9lurxiQiIW209GnZEjX0=";
  42     };
  43
  44     patches = [
  45       # Convert np.int -> np.int32 since np.int is deprecated, https://github.com/google-deepmind/bsuite/pull/48
  46       (fetchpatch {
  47         url = "https://github.com/google-deepmind/bsuite/pull/48/commits/f8d81b2f1c27ef2c8c71ae286001ed879ea306ab.patch";
  48         hash = "sha256-FXtvVS+U8brulq8Z27+yWIimB+kigGiUOIv1SHb1TA8=";
  49       })
  50       # Replace imp with importlib, https://github.com/google-deepmind/bsuite/pull/50
  51       (fetchpatch {
  52         name = "replace-imp.patch";
  53         url = "https://github.com/google-deepmind/bsuite/commit/d08b63655c7efa5b5bb0f35e825e17549d23e812.patch";
  54         hash = "sha256-V5p/6edNXTpEckuSuxJ/mvfJng5yE/pfeMoYbvlNpEo=";
  55       })
  56     ];
  57
  58     build-system = [ setuptools ];
  59
  60     dependencies = [
  61       absl-py
  62       dm-env
  63       dm-tree
  64       frozendict
  65       gym
  66       matplotlib
  67       mizani
  68       pandas
  69       patsy
  70       plotnine
  71       scikit-image
  72       scipy
  73       statsmodels
  74       termcolor
  75     ];
  76
  77     nativeCheckInputs = [
  78       distrax
  79       dm-haiku
  80       dm-sonnet
  81       optax
  82       pytestCheckHook
  83       rlax
  84       tensorflow-probability
  85     ];
  86
  87     pythonImportsCheck = [ "bsuite" ];
  88
  89     disabledTests = [
  90       # Tests require network connection
  91       "test_run9"
  92       "test_longer_action_sequence"
  93       "test_reset"
  94       "test_step_after_reset"
  95       "test_step_on_fresh_environment"
  96       "test_longer_action_sequence"
  97       "test_reset"
  98       "test_step_after_reset"
  99       "test_step_on_fresh_environment"
 100       "test_longer_action_sequence"
 101       "test_reset"
 102       "test_step_after_reset"
 103       "test_step_on_fresh_environment"
 104       "test_logger"
 105       "test_episode_truncation"
 106     ];
 107
 108     # Escape infinite recursion with rlax
 109     doCheck = false;
 110
 111     passthru.tests = {
 112       check = bsuite.overridePythonAttrs (_: {
 113         doCheck = true;
 114       });
 115     };
 116
 117     meta = with lib; {
 118       description = "Collection of experiments that investigate core capabilities of a reinforcement learning (RL) agent";
 119       homepage = "https://github.com/deepmind/bsuite";
 120       changelog = "https://github.com/google-deepmind/bsuite/releases/tag/${version}";
 121       license = licenses.asl20;
 122       maintainers = with maintainers; [ onny ];
 123     };
 124   };
 125 in
 126 bsuite