1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
|
{ lib
, buildPythonPackage
, pythonOlder
, fetchFromGitHub
, ninja
, setuptools
, wheel
, which
, cloudpickle
, numpy
, torch
, ale-py
, gym
, pygame
, torchsnapshot
, gymnasium
, mujoco
, h5py
, huggingface-hub
, minari
, pandas
, pillow
, requests
, scikit-learn
, torchvision
, tqdm
, moviepy
, git
, hydra-core
, tensorboard
, wandb
, packaging
, tensordict
, imageio
, pytest-rerunfailures
, pytestCheckHook
, pyyaml
, scipy
, stdenv
}:
buildPythonPackage rec {
pname = "torchrl";
version = "0.4.0";
pyproject = true;
disabled = pythonOlder "3.8";
src = fetchFromGitHub {
owner = "pytorch";
repo = "rl";
rev = "refs/tags/v${version}";
hash = "sha256-8wSyyErqveP9zZS/UGvWVBYyylu9BuA447GEjXIzBIk=";
};
build-system = [
ninja
setuptools
wheel
which
];
dependencies = [
cloudpickle
numpy
packaging
tensordict
torch
];
passthru.optional-dependencies = {
atari = [
ale-py
gym
pygame
];
checkpointing = [
torchsnapshot
];
gym-continuous = [
gymnasium
mujoco
];
offline-data = [
h5py
huggingface-hub
minari
pandas
pillow
requests
scikit-learn
torchvision
tqdm
];
rendering = [
moviepy
];
utils = [
git
hydra-core
tensorboard
tqdm
wandb
];
};
# torchrl needs to create a folder to store datasets
preBuild = ''
export D4RL_DATASET_DIR=$(mktemp -d)
'';
pythonImportsCheck = [
"torchrl"
];
# We have to delete the source because otherwise it is used instead of the installed package.
preCheck = ''
rm -rf torchrl
export XDG_RUNTIME_DIR=$(mktemp -d)
'';
nativeCheckInputs = [
gymnasium
imageio
pytest-rerunfailures
pytestCheckHook
pyyaml
scipy
torchvision
]
++ passthru.optional-dependencies.atari
++ passthru.optional-dependencies.gym-continuous
++ passthru.optional-dependencies.rendering;
disabledTests = [
# mujoco.FatalError: an OpenGL platform library has not been loaded into this process, this most likely means that a valid OpenGL context has not been created before mjr_makeContext was called
"test_vecenvs_env"
# ValueError: Can't write images with one color channel.
"test_log_video"
# Those tests require the ALE environments (provided by unpackaged shimmy)
"test_collector_env_reset"
"test_gym"
"test_gym_fake_td"
"test_recorder"
"test_recorder_load"
"test_rollout"
"test_parallel_trans_env_check"
"test_serial_trans_env_check"
"test_single_trans_env_check"
"test_td_creation_from_spec"
"test_trans_parallel_env_check"
"test_trans_serial_env_check"
"test_transform_env"
];
meta = with lib; {
description = "A modular, primitive-first, python-first PyTorch library for Reinforcement Learning";
homepage = "https://github.com/pytorch/rl";
changelog = "https://github.com/pytorch/rl/releases/tag/v${version}";
license = licenses.mit;
maintainers = with maintainers; [ GaetanLepage ];
# ~3k tests fail with: RuntimeError: internal error
broken = stdenv.isLinux && stdenv.isAarch64;
};
}
|
