diff --git a/pkgs/development/python-modules/tensordict/default.nix b/pkgs/development/python-modules/tensordict/default.nix index 317807f2e6fc..0d3938dd8775 100644 --- a/pkgs/development/python-modules/tensordict/default.nix +++ b/pkgs/development/python-modules/tensordict/default.nix @@ -29,16 +29,21 @@ buildPythonPackage rec { pname = "tensordict"; - version = "0.9.1"; + version = "0.10.0"; pyproject = true; src = fetchFromGitHub { owner = "pytorch"; repo = "tensordict"; tag = "v${version}"; - hash = "sha256-OdS9dw/BtSLZuY857O2njlFOMQj5IJ6v9c2aRP+H1Hc="; + hash = "sha256-yxyA9BfN2hp1C3s+g2zBM2gVtckH3LV7luWw8DshFUs="; }; + postPatch = '' + substituteInPlace pyproject.toml \ + --replace-fail "pybind11[global]" "pybind11" + ''; + build-system = [ pybind11 setuptools diff --git a/pkgs/development/python-modules/torchrl/default.nix b/pkgs/development/python-modules/torchrl/default.nix index b4f2d48a021a..891d98f35a3c 100644 --- a/pkgs/development/python-modules/torchrl/default.nix +++ b/pkgs/development/python-modules/torchrl/default.nix @@ -5,22 +5,29 @@ fetchFromGitHub, # build-system + cmake, ninja, + numpy, + pybind11, setuptools, - which, + torch, # dependencies cloudpickle, - numpy, packaging, + pyvers, tensordict, - torch, # optional-dependencies # atari gymnasium, + # brax + brax, + jax, # checkpointing torchsnapshot, + # dm-control + dm-control, # gym-continuous mujoco, # llm @@ -36,6 +43,8 @@ sentencepiece, transformers, vllm, + # marl + pettingzoo, # offline-data h5py, huggingface-hub, @@ -64,33 +73,51 @@ buildPythonPackage rec { pname = "torchrl"; - version = "0.9.2"; + version = "0.10.0"; pyproject = true; src = fetchFromGitHub { owner = "pytorch"; repo = "rl"; tag = "v${version}"; - hash = "sha256-6rU5+J70T0E7+60jihsjwlLls8jJlxKi3nmrL0xm2c0="; + hash = "sha256-DqLB1JnQ96cxVEzcXra1hFVfrN7eXTlTwPtlPClnaBA="; }; + postPatch = '' + substituteInPlace pyproject.toml \ + --replace-fail "pybind11[global]" "pybind11" + ''; + build-system = [ + cmake ninja + numpy + pybind11 setuptools - which + torch ]; + dontUseCmakeConfigure = true; dependencies = [ cloudpickle numpy packaging tensordict + pyvers torch ]; optional-dependencies = { - atari = gymnasium.optional-dependencies.atari; + atari = [ + gymnasium + ] + ++ gymnasium.optional-dependencies.atari; + brax = [ + brax + jax + ]; checkpointing = [ torchsnapshot ]; + dm-control = [ dm-control ]; gym-continuous = [ gymnasium mujoco @@ -109,6 +136,11 @@ buildPythonPackage rec { transformers vllm ]; + marl = [ + # dm-meltingpot (unpackaged) + pettingzoo + # vmas (unpackaged) + ]; offline-data = [ h5py huggingface-hub @@ -120,10 +152,15 @@ buildPythonPackage rec { torchvision tqdm ]; + open-spiel = [ + # open-spiel (unpackaged) + ]; rendering = [ moviepy ]; + replay-buffer = [ torch ]; utils = [ git hydra-core + # hydra-submitit-launcher (unpackaged) tensorboard tqdm wandb