pyrox.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.6.0";
26 pyproject = true;
27
28 src = fetchFromGitHub {
29 owner = "DLR-RM";
30 repo = "stable-baselines3";
31 tag = "v${version}";
32 hash = "sha256-VnoQ8cKqPcZPpR9c3M6xJDdG7gnO9fxIa4v2kxd9Nzg=";
33 };
34
35 build-system = [ setuptools ];
36
37 pythonRelaxDeps = [
38 "gymnasium"
39 ];
40
41 dependencies = [
42 cloudpickle
43 gymnasium
44 matplotlib
45 numpy
46 pandas
47 torch
48 ];
49
50 nativeCheckInputs = [
51 ale-py
52 pytestCheckHook
53 rich
54 tqdm
55 ];
56
57 pythonImportsCheck = [ "stable_baselines3" ];
58
59 disabledTestPaths = [
60 # Tests starts training a model, which takes too long
61 "tests/test_cnn.py"
62 "tests/test_dict_env.py"
63 "tests/test_her.py"
64 "tests/test_save_load.py"
65 ];
66
67 disabledTests = [
68 # Flaky: Can fail if it takes too long, which happens when the system is under heavy load
69 "test_fps_logger"
70
71 # Tests that attempt to access the filesystem
72 "test_make_atari_env"
73 "test_vec_env_monitor_kwargs"
74 ];
75
76 meta = {
77 description = "PyTorch version of Stable Baselines, reliable implementations of reinforcement learning algorithms";
78 homepage = "https://github.com/DLR-RM/stable-baselines3";
79 changelog = "https://github.com/DLR-RM/stable-baselines3/releases/tag/v${version}";
80 license = lib.licenses.mit;
81 maintainers = with lib.maintainers; [ derdennisop ];
82 };
83}