tjh.dev
1{
2 lib,
3 buildPythonPackage,
4 fetchFromGitHub,
5
6 # build-system
7 setuptools,
8
9 # dependencies
10 cloudpickle,
11 gymnasium,
12 matplotlib,
13 numpy,
14 pandas,
15 torch,
16
17 # tests
18 ale-py,
19 pytestCheckHook,
20 rich,
21 tqdm,
22}:
23buildPythonPackage rec {
24 pname = "stable-baselines3";
25 version = "2.4.0";
26 pyproject = true;
27
28 src = fetchFromGitHub {
29 owner = "DLR-RM";
30 repo = "stable-baselines3";
31 rev = "refs/tags/v${version}";
32 hash = "sha256-OFmjAkUS0wrns5CkWjqR7zTjKLVPI/NrpVIKJhLAwYM=";
33 };
34
35 build-system = [ setuptools ];
36
37 dependencies = [
38 cloudpickle
39 gymnasium
40 matplotlib
41 numpy
42 pandas
43 torch
44 ];
45
46 nativeCheckInputs = [
47 ale-py
48 pytestCheckHook
49 rich
50 tqdm
51 ];
52
53 pythonImportsCheck = [ "stable_baselines3" ];
54
55 disabledTestPaths = [
56 # Tests starts training a model, which takes too long
57 "tests/test_cnn.py"
58 "tests/test_dict_env.py"
59 "tests/test_her.py"
60 "tests/test_save_load.py"
61 ];
62
63 disabledTests = [
64 # Tests that attempt to access the filesystem
65 "test_make_atari_env"
66 "test_vec_env_monitor_kwargs"
67 ];
68
69 meta = {
70 description = "PyTorch version of Stable Baselines, reliable implementations of reinforcement learning algorithms";
71 homepage = "https://github.com/DLR-RM/stable-baselines3";
72 changelog = "https://github.com/DLR-RM/stable-baselines3/releases/tag/v${version}";
73 license = lib.licenses.mit;
74 maintainers = with lib.maintainers; [ derdennisop ];
75 };
76}