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
|
{ lib
, buildPythonPackage
, pythonOlder
, fetchFromGitHub
, ninja
, setuptools
, wheel
, which
, cloudpickle
, numpy
, torch
, ale-py
, gym
, pygame
, gymnasium
, mujoco
, moviepy
, git
, hydra-core
, tensorboard
, tqdm
, wandb
, packaging
, tensordict
, imageio
, pytest-rerunfailures
, pytestCheckHook
, pyyaml
, scipy
}:
buildPythonPackage rec {
pname = "torchrl";
version = "0.3.1";
pyproject = true;
disabled = pythonOlder "3.8";
src = fetchFromGitHub {
owner = "pytorch";
repo = "rl";
rev = "refs/tags/v${version}";
hash = "sha256-lETW996IKPUGgZpe+cyzrXvVmDSwj5G4XFreFmGxReQ=";
};
nativeBuildInputs = [
ninja
setuptools
wheel
which
];
propagatedBuildInputs = [
cloudpickle
numpy
packaging
tensordict
torch
];
passthru.optional-dependencies = {
atari = [
ale-py
gym
pygame
];
gym-continuous = [
gymnasium
mujoco
];
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
]
++ 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 ];
};
}
|
