Merge pull request #309919 from GaetanLepage/torchrl

python311Packages.torchrl: 0.3.1 -> 0.4.0

pkgs/development/python-modules/tensordict/default.nix

@@ -11,11 +11,12 @@
 , h5py
 , pytestCheckHook
 , stdenv
+, pythonAtLeast
 }:
 
 buildPythonPackage rec {
   pname = "tensordict";
-  version = "0.3.1";
+  version = "0.4.0";
   pyproject = true;
 
   disabled = pythonOlder "3.8";
@@ -24,17 +25,17 @@ buildPythonPackage rec {
     owner = "pytorch";
     repo = "tensordict";
     rev = "refs/tags/v${version}";
-    hash = "sha256-eCx1r7goqOdGX/0mSGCiLhdGQTh4Swa5aFiLSsL56p0=";
+    hash = "sha256-wKEzNaaazGEkoElzp93RIlq/r5uRUdM7UyDy/DygIEc=";
   };
 
-  nativeBuildInputs = [
+  build-system = [
     setuptools
     torch
     wheel
     which
   ];
 
-  propagatedBuildInputs = [
+  dependencies = [
     cloudpickle
     numpy
     torch
@@ -72,5 +73,7 @@ buildPythonPackage rec {
     homepage = "https://github.com/pytorch/tensordict";
     license = licenses.mit;
     maintainers = with maintainers; [ GaetanLepage ];
+    # No python 3.12 support yet: https://github.com/pytorch/rl/issues/2035
+    broken = pythonAtLeast "3.12";
   };
 }

pkgs/development/python-modules/torchrl/default.nix

@@ -12,13 +12,22 @@
 , ale-py
 , gym
 , pygame
+, torchsnapshot
 , gymnasium
 , mujoco
+, h5py
+, huggingface-hub
+, minari
+, pandas
+, pillow
+, requests
+, scikit-learn
+, torchvision
+, tqdm
 , moviepy
 , git
 , hydra-core
 , tensorboard
-, tqdm
 , wandb
 , packaging
 , tensordict
@@ -31,7 +40,7 @@
 
 buildPythonPackage rec {
   pname = "torchrl";
-  version = "0.3.1";
+  version = "0.4.0";
   pyproject = true;
 
   disabled = pythonOlder "3.8";
@@ -40,17 +49,17 @@ buildPythonPackage rec {
     owner = "pytorch";
     repo = "rl";
     rev = "refs/tags/v${version}";
-    hash = "sha256-lETW996IKPUGgZpe+cyzrXvVmDSwj5G4XFreFmGxReQ=";
+    hash = "sha256-8wSyyErqveP9zZS/UGvWVBYyylu9BuA447GEjXIzBIk=";
   };
 
-  nativeBuildInputs = [
+  build-system = [
     ninja
     setuptools
     wheel
     which
   ];
 
-  propagatedBuildInputs = [
+  dependencies = [
     cloudpickle
     numpy
     packaging
@@ -64,10 +73,24 @@ buildPythonPackage rec {
       gym
       pygame
     ];
+    checkpointing = [
+      torchsnapshot
+    ];
     gym-continuous = [
       gymnasium
       mujoco
     ];
+    offline-data = [
+      h5py
+      huggingface-hub
+      minari
+      pandas
+      pillow
+      requests
+      scikit-learn
+      torchvision
+      tqdm
+    ];
     rendering = [
       moviepy
     ];
@@ -103,6 +126,7 @@ buildPythonPackage rec {
     pytestCheckHook
     pyyaml
     scipy
+    torchvision
   ]
   ++ passthru.optional-dependencies.atari
   ++ passthru.optional-dependencies.gym-continuous