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