tianshou: init at 1.1.0; sensai-utils: init at 1.4.0; rlcard: patch distutils; pettingzoo: fix dependencies (#384526)

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

@@ -1,27 +1,28 @@
 {
   lib,
   buildPythonPackage,
-  pythonOlder,
+  chess,
   fetchFromGitHub,
-  setuptools,
-  wheel,
   gymnasium,
   numpy,
-  chess,
   pillow,
+  pre-commit,
   pybox2d,
   pygame,
   pymunk,
-  rlcard,
-  scipy,
-  pre-commit,
   pynput,
   pytest,
   pytest-cov-stub,
   pytest-markdown-docs,
   pytest-xdist,
   pytestCheckHook,
+  pythonOlder,
+  rlcard,
+  scipy,
+  setuptools,
+  shimmy,
   stdenv,
+  wheel,
 }:
 
 buildPythonPackage rec {
@@ -49,17 +50,7 @@ buildPythonPackage rec {
   ];
 
   optional-dependencies = {
-    all = [
-      chess
-      # multi-agent-ale-py
-      pillow
-      pybox2d
-      pygame
-      pymunk
-      rlcard
-      scipy
-      # shimmy
-    ];
+    all = lib.flatten (lib.attrValues (lib.filterAttrs (n: v: n != "all") optional-dependencies));
     atari = [
       # multi-agent-ale-py
       pygame
@@ -72,7 +63,7 @@ buildPythonPackage rec {
       chess
       pygame
       rlcard
-      # shimmy
+      shimmy
     ];
     mpe = [ pygame ];
     other = [ pillow ];

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

@@ -9,7 +9,7 @@
   termcolor,
   pytestCheckHook,
   torch,
-  pythonAtLeast,
+  fetchpatch2,
 }:
 
 buildPythonPackage rec {
@@ -24,6 +24,16 @@ buildPythonPackage rec {
     hash = "sha256-SWj6DBItQzSM+nioV54a350Li7tbBaVXsQxNAqVgB0k=";
   };
 
+  patches = [
+    # Remove distutils to make it compatible with Python 3.12
+    # https://github.com/datamllab/rlcard/pull/323
+    (fetchpatch2 {
+      name = "remove-distutils.patch";
+      url = "https://github.com/datamllab/rlcard/commit/e44378157aaf229ffe2aaef9fafe500c2844045e.patch";
+      hash = "sha256-aQS4d9ETj6pDv26G77mC+0xHQMA2hjspAxtAyz0rA6Y=";
+    })
+  ];
+
   build-system = [
     setuptools
     wheel
@@ -70,7 +80,5 @@ buildPythonPackage rec {
     changelog = "https://github.com/datamllab/rlcard/releases/tag/${version}";
     license = licenses.mit;
     maintainers = with maintainers; [ GaetanLepage ];
-    # Relies on deprecated distutils
-    broken = pythonAtLeast "3.12";
   };
 }

pkgs/development/python-modules/sensai-utils/default.nix

@@ -0,0 +1,40 @@
+{
+  lib,
+  buildPythonPackage,
+  fetchFromGitHub,
+  pytestCheckHook,
+  pythonOlder,
+  setuptools,
+  typing-extensions,
+}:
+
+buildPythonPackage rec {
+  pname = "sensai-utils";
+  version = "1.4.0";
+  pyproject = true;
+
+  disabled = pythonOlder "3.11";
+
+  src = fetchFromGitHub {
+    owner = "opcode81";
+    repo = "sensAI-utils";
+    tag = "v${version}";
+    hash = "sha256-XgZv76tLeTRCvNptasp8EiU2DC+HWkc1xhlCA+YiUZY=";
+  };
+
+  build-system = [ setuptools ];
+
+  dependencies = [ typing-extensions ];
+
+  pythonImportsCheck = [ "sensai.util" ];
+
+  nativeCheckInputs = [ pytestCheckHook ];
+
+  meta = with lib; {
+    description = "Utilities from sensAI, the Python library for sensible AI";
+    homepage = "https://github.com/opcode81/sensAI-utils";
+    changelog = "https://github.com/opcode81/sensAI-utils/releases/tag/v${version}";
+    license = licenses.mit;
+    maintainers = with maintainers; [ derdennisop ];
+  };
+}

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

@@ -0,0 +1,175 @@
+{
+  lib,
+  ale-py,
+  buildPythonPackage,
+  cython,
+  deepdiff,
+  docstring-parser,
+  fetchFromGitHub,
+  gymnasium,
+  h5py,
+  imageio,
+  joblib,
+  jsonargparse,
+  matplotlib,
+  mujoco,
+  numba,
+  numpy,
+  opencv,
+  overrides,
+  packaging,
+  pandas,
+  pettingzoo,
+  poetry-core,
+  pybox2d,
+  pybullet,
+  pygame,
+  pymunk,
+  pytestCheckHook,
+  pythonOlder,
+  scipy,
+  sensai-utils,
+  shimmy,
+  swig,
+  tensorboard,
+  torch,
+  tqdm,
+}:
+
+buildPythonPackage rec {
+  pname = "tianshou";
+  version = "1.1.0";
+  pyproject = true;
+
+  disabled = pythonOlder "3.11";
+
+  src = fetchFromGitHub {
+    owner = "thu-ml";
+    repo = "tianshou";
+    tag = "v${version}";
+    hash = "sha256-eiwbSX8Q3KF6h7CfjuZ+7HlXwpvLga1NVr3e+FkPaHc=";
+  };
+
+  pythonRelaxDeps = [
+    "deepdiff"
+    "gymnasium"
+    "numpy"
+  ];
+
+  pythonRemoveDeps = [ "virtualenv" ];
+
+  postPatch = ''
+    # silence matplotlib warning
+    export MPLCONFIGDIR=$(mktemp -d)
+  '';
+
+  build-system = [ poetry-core ];
+
+  dependencies = [
+    deepdiff
+    gymnasium
+    h5py
+    matplotlib
+    numba
+    numpy
+    overrides
+    packaging
+    pandas
+    pettingzoo
+    sensai-utils
+    tensorboard
+    torch
+    tqdm
+  ];
+
+  optional-dependencies = {
+    all = lib.flatten (lib.attrValues (lib.filterAttrs (n: v: n != "all") optional-dependencies));
+
+    argparse = [
+      docstring-parser
+      jsonargparse
+    ];
+
+    atari = [
+      ale-py
+      # autorom
+      opencv
+      shimmy
+    ];
+
+    box2d = [
+      # instead of box2d-py
+      pybox2d
+      pygame
+      swig
+    ];
+
+    classic_control = [
+      pygame
+    ];
+
+    mujoco = [
+      mujoco
+      imageio
+      cython
+    ];
+
+    pybullet = [
+      pybullet
+    ];
+
+    # envpool = [
+    #   envpool
+    # ];
+
+    # robotics = [
+    #   gymnasium-robotics
+    # ];
+
+    # vizdoom = [
+    #   vizdoom
+    # ];
+
+    eval = [
+      docstring-parser
+      joblib
+      jsonargparse
+      # rliable
+      scipy
+    ];
+  };
+
+  pythonImportsCheck = [ "tianshou" ];
+
+  nativeCheckInputs = [
+    pygame
+    pymunk
+    pytestCheckHook
+  ];
+
+  disabledTestPaths = [
+    # remove tests that require lot of compute (ai model training tests)
+    "test/continuous"
+    "test/discrete"
+    "test/highlevel"
+    "test/modelbased"
+    "test/offline"
+  ];
+
+  disabledTests = [
+    # AttributeError: 'TimeLimit' object has no attribute 'test_attribute'
+    "test_attr_unwrapped"
+    # Failed: DID NOT RAISE <class 'TypeError'>
+    "test_batch"
+    # Failed: Raised AssertionError
+    "test_vecenv"
+  ];
+
+  meta = with lib; {
+    description = "Elegant PyTorch deep reinforcement learning library";
+    homepage = "https://github.com/thu-ml/tianshou";
+    changelog = "https://github.com/thu-ml/tianshou/releases/tag/v${version}";
+    license = licenses.mit;
+    maintainers = with maintainers; [ derdennisop ];
+  };
+}

pkgs/top-level/python-packages.nix

@@ -14864,6 +14864,8 @@ self: super: with self; {
 
   senf = callPackage ../development/python-modules/senf { };
 
+  sensai-utils = callPackage ../development/python-modules/sensai-utils { };
+
   sense-energy = callPackage ../development/python-modules/sense-energy { };
 
   sensirion-ble = callPackage ../development/python-modules/sensirion-ble { };
@@ -16422,6 +16424,8 @@ self: super: with self; {
     py = python.override (lib.optionalAttrs (!python.isPyPy) { x11Support = true; });
   };
 
+  tianshou = callPackage ../development/python-modules/tianshou { };
+
   tidalapi = callPackage ../development/python-modules/tidalapi { };
 
   tidyexc = callPackage ../development/python-modules/tidyexc { };