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
31let
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 };
125in
126bsuite