DRAFT: add maniskill3 package along with all of its dependencies

flake.nix

@@ -3,7 +3,6 @@
 
   inputs = {
     nixpkgs.url = "github:NixOS/nixpkgs/nixos-24.05";
-
     utils.url = "github:numtide/flake-utils";
   };
 
@@ -14,7 +13,8 @@
       torch-family = import ./overlays/torch-family.nix;
       jax-family = import ./overlays/jax-family.nix;
       data-utils = import ./overlays/data-utils.nix;
-      simulators = import ./overlays/simulators.nix;
+      simulators = import ./overlays/haosu.nix;
+      haosu = import ./overlays/simulators.nix;
       math = import ./overlays/math.nix;
       misc = import ./overlays/misc.nix;
       apis = import ./overlays/apis.nix;
@@ -29,6 +29,7 @@
         self.overlays.jax-family
         self.overlays.data-utils
         self.overlays.simulators
+        self.overlays.haosu
         self.overlays.math
         self.overlays.misc
         self.overlays.apis
@@ -70,10 +71,9 @@
         jax = pkgs.callPackage ./pkgs/dev-shell/jax.nix {};
         quicktest = pkgs.callPackage ./pkgs/dev-shell/quicktest.nix {};
         ddddocr = pkgs.callPackage ./pkgs/dev-shell/ddddocr.nix {};
+        maniskill = pkgs.callPackage ./pkgs/dev-shell/maniskill.nix {};
+        habitat = pkgs.callPackage ./pkgs/dev-shell/habitat.nix {};
       };
-      # devShells.py38 = pkgs.callPackage ./pkgs/dev-shell {
-      #   python3 = pkgs.python38;
-      # };
 
       packages = {
         inherit (pkgs.python3Packages)
@@ -102,7 +102,6 @@
           mujoco-menagerie
           dm-control
           python-fcl
-          sapien
           glfw
           gputil
           aws-requests-auth
@@ -113,6 +112,21 @@
           robosuite
           robocasa
 
+          sapien
+          maniskill
+          pytorch-seed
+          pytorch-kinematics
+          arm-pytorch-utilities
+          fast-kinematics
+          tyro
+          toppra
+          mplib
+          stable-baselines3
+
+          magnum-bindings
+          habitat-sim
+          habitat-lab
+
           # ----- Math -----
           numpy-quaternion
           chumpy
@@ -153,7 +167,10 @@
           quickcpplib
           status-code
           eve
-          aider-chat;
+          physx5
+          physx5-gpu
+          sapien-vulkan-2
+          ompl;
       };
 
       apps = {

overlays/haosu.nix

@@ -0,0 +1,75 @@
+final: prev: {
+  pythonPackagesExtensions = prev.pythonPackagesExtensions ++ [
+    (python-final: python-prev: {
+      sapien = python-final.callPackage ../pkgs/sapien {};
+      maniskill = python-final.callPackage ../pkgs/maniskill {};
+
+      pytorch-seed = python-final.callPackage ../pkgs/pytorch-seed {
+        torch = python-final.torchWithCuda;
+      };
+      pytorch-kinematics = python-final.callPackage ../pkgs/pytorch-kinematics {
+        torch = python-final.torchWithCuda;
+      };
+      arm-pytorch-utilities = python-final.callPackage ../pkgs/arm-pytorch-utilities {
+        torch = python-final.torchWithCuda;
+      };
+      fast-kinematics = python-final.callPackage ../pkgs/fast-kinematics {
+        torch = python-final.torchWithCuda;
+      };
+      toppra = python-final.callPackage ../pkgs/toppra {};
+      mplib = python-final.callPackage ../pkgs/mplib {};
+      stable-baselines3 = python-final.callPackage ../pkgs/stable-baselines3 {};
+
+      tyro = python-final.callPackage ../pkgs/tyro {};
+      docstring-parser = python-prev.docstring-parser.overrideAttrs rec {
+        version = "0.16";
+        src = final.fetchFromGitHub {
+          owner = "rr-";
+          repo = "docstring_parser";
+          rev = "${version}";
+          hash = "sha256-xwV+mgCOC/MyCqGELkJVqQ3p2g2yw/Ieomc7k0HMXms=";
+        };
+      };
+
+      pybind-smart-holder = python-final.pybind11.overrideAttrs {
+        src = final.fetchFromGitHub {
+          owner = "pybind";
+          repo = "pybind11";
+          rev = "3b35ce475fa359abc31d979972f650c601d6158b";
+          hash = "sha256-zgWTcgO0BvCOjFrNPrLTi0JXedhW2Oai1qwf5DA7e6A=";
+        };
+        postPatch = "";
+      };
+    })
+  ];
+
+  physx5-gpu = final.callPackage ../pkgs/physx5-gpu {};
+  physx5 = final.callPackage ../pkgs/physx5 {};
+  sapien-vulkan-2 = final.callPackage ../pkgs/sapien-vulkan-2 {};
+  assimp-sapien = final.assimp.overrideAttrs {
+    src = final.fetchFromGitHub {
+      owner = "fbxiang";
+      repo = "assimp";
+      rev = "0ea31aa6734336dc1e62c6d9bde3e49b6d71b811";
+      sha256 = "sha256-IqF46UQNGQ/EZJ/D0SsOqp+Tyn5oSNtunNx0lxaTRGE=";
+    };
+  };
+  glm-sapien = final.glm.overrideAttrs {
+    src = final.fetchFromGitHub {
+      owner = "g-truc";
+      repo = "glm";
+      rev = "0.9.9.8";
+      hash = "sha256-F//+3L5Ozrw6s7t4LrcUmO7sN30ZSESdrPAYX57zgr8=";
+    };
+  };
+  openexr-sapien = final.openexr_3.overrideAttrs {
+    postFixup = ''
+      substituteInPlace $dev/include/OpenEXR/*.h \
+        --replace-warn '#include <Imath' '#include <Imath/Imath' \
+        --replace-warn '#include "Imath' '#include "Imath/Imath' \
+        --replace-warn '#include <half.h' '#include <Imath/half.h'
+    '';
+    doCheck = false;
+  };
+  ompl = final.callPackage ../pkgs/ompl {};
+}

overlays/jax-family.nix

@@ -1,7 +1,7 @@
 final: prev: {
   pythonPackagesExtensions = prev.pythonPackagesExtensions ++ [
     (python-final: python-prev: {
-      jaxlib = python-prev.jaxlib.override {
+      jaxlib = python-final.callPackage ../bleeding/jaxlib {
         inherit (final.config) cudaSupport;
       };
 

overlays/simulators.nix

@@ -36,9 +36,6 @@ final: prev: {
       mujoco-menagerie = python-final.callPackage ../pkgs/mujoco-menagerie {};
       python-fcl = python-final.callPackage ../pkgs/python-fcl {};
 
-      # SAPIEN from Hao Su's Lab
-      sapien = python-final.callPackage ../pkgs/sapien {};
-
       # nixpkgs has a problematic glfw on 23.11 at this moment.
       glfw = python-final.callPackage ../bleeding/glfw {};
 
@@ -53,9 +50,19 @@ final: prev: {
       # Robocasa/Robosuite
       robosuite = python-final.callPackage ../research-community/robosuite {};
       robocasa = python-final.callPackage ../research-community/robocasa {};
+
+      habitat-sim = python-final.callPackage ../pkgs/habitat-sim {};
+      habitat-lab = python-final.callPackage ../pkgs/habitat-lab {};
+      magnum-bindings = python-final.callPackage ../pkgs/magnum-bindings {};
     })
   ];
 
   # mujoco depends on sdflib now, which in turn depends on cereal >= 1.3.1
   cereal = final.cereal_1_3_2;
+  magnum = final.callPackage ../pkgs/magnum {};
+  magnum-integration = final.callPackage ../pkgs/magnum-integration {};
+  magnum-plugins = final.callPackage ../pkgs/magnum-plugins {};
+  corrade = final.callPackage ../pkgs/corrade {};
+  basis-universal = final.callPackage ../pkgs/basis-universal {};
+  recastnavigation = final.callPackage ../pkgs/recastnavigation {};
 }

overlays/torch-family.nix

@@ -12,20 +12,20 @@ final: prev:
       # https://github.com/SomeoneSerge/dust3r.nix/commit/8715f78a6f875fa5a1795da8f56e05fe2b2811d6
       #
       # The proper fix is still a pending PR: https://github.com/NixOS/nixpkgs/pull/297590
-      torchWithCuda = python-prev.torchWithCuda.overridePythonAttrs (oldAttrs: {
-        extraRunpaths = [ "${prev.lib.getLib final.cudaPackages.cuda_nvrtc}/lib" ];
-        postPhases = prev.lib.optionals final.stdenv.hostPlatform.isUnix ["postPatchelfPhase" ];
-        postPatchelfPhase = ''
-          while IFS= read -r -d $'\0' elf ; do
-            for extra in $extraRunpaths ; do
-              echo patchelf "$elf" --add-rpath "$extra" >&2
-              patchelf "$elf" --add-rpath "$extra"
-            done
-          done < <(
-            find "''${!outputLib}" "$out" -type f -iname '*.so' -print0
-          )
-        '';
-      });
+      # torchWithCuda = python-prev.torchWithCuda.overridePythonAttrs (oldAttrs: {
+      #   extraRunpaths = [ "${prev.lib.getLib final.cudaPackages.cuda_nvrtc}/lib" ];
+      #   postPhases = prev.lib.optionals final.stdenv.hostPlatform.isUnix ["postPatchelfPhase" ];
+      #   postPatchelfPhase = ''
+      #     while IFS= read -r -d $'\0' elf ; do
+      #       for extra in $extraRunpaths ; do
+      #         echo patchelf "$elf" --add-rpath "$extra" >&2
+      #         patchelf "$elf" --add-rpath "$extra"
+      #       done
+      #     done < <(
+      #       find "''${!outputLib}" "$out" -type f -iname '*.so' -print0
+      #     )
+      #   '';
+      # });
 
       LIV-robotics = python-final.callPackage ../pkgs/LIV-robotics {
         pytorch = python-final.torchWithCuda;

pkgs/arm-pytorch-utilities/default.nix

@@ -0,0 +1,45 @@
+{ lib
+, pythonOlder
+, buildPythonPackage
+, fetchFromGitHub
+, setuptools
+, numpy
+, torch
+, matplotlib
+, pytorch-seed
+, scipy
+}:
+
+buildPythonPackage rec {
+  pname = "arm_pytorch_utilities";
+  version = "0.4.3";
+  format = "pyproject";
+
+  src = fetchFromGitHub {
+    owner = "UM-ARM-Lab";
+    repo = pname;
+    rev = "6312d5094f71a646041fa2994663a19f052ee857";
+    hash = "sha256-CTe/cFc2KKG56IeLDDZiDt6GgBGe/vA/usqilvigC80=";
+  };
+
+  buildInputs = [
+    setuptools
+  ];
+
+  propagatedBuildInputs = [
+    numpy
+    torch
+    matplotlib
+    pytorch-seed
+    scipy
+  ];
+
+  pythonImportsCheck = [ "arm_pytorch_utilities" ];
+
+  meta = with lib; {
+    description = "University of Michigan ARM Lab PyTorch utilities";
+    homepage = "https://github.com/UM-ARM-Lab/arm_pytorch_utilities";
+    license = licenses.mit;
+    maintainers = with maintainers; [ breakds ];
+  };
+}

pkgs/basis-universal/default.nix

@@ -0,0 +1,31 @@
+{
+  lib,
+  stdenv,
+  fetchFromGitHub,
+  cmake,
+}:
+
+stdenv.mkDerivation rec {
+  pname = "basis-universal";
+  version = "1_15_update2";
+
+  # Breaks magnum-plugins with every update, follow
+  # https://github.com/mosra/archlinux/blob/36275b5aea5e1521735b5da7819aea9ce1300976/basis-universal-src/PKGBUILD
+  src = fetchFromGitHub {
+    owner = "BinomialLLC";
+    repo = "basis_universal";
+    rev = "v${version}";
+    hash = "sha256-2snzq/SnhWHIgSbUUgh24B6tka7EfkGO+nwKEObRkU4=";
+  };
+
+  nativeBuildInputs = [ cmake ];
+
+  meta = with lib; {
+    description = "Basis Universal GPU Texture Codec";
+    homepage = "https://github.com/BinomialLLC/basis_universal";
+    license = licenses.asl20;
+    maintainers = with maintainers; [ ];
+    mainProgram = "basis-universal";
+    platforms = platforms.all;
+  };
+}

pkgs/corrade/default.nix

@@ -0,0 +1,32 @@
+{
+  lib,
+  stdenv,
+  fetchFromGitHub,
+  cmake,
+}:
+
+stdenv.mkDerivation rec {
+  pname = "corrade";
+  version = "unstable-2024-03-05";
+
+  src = fetchFromGitHub {
+    owner = "mosra";
+    repo = "corrade";
+    rev = "295bbba1f49887da060465f88b8501965f6acd7d";
+    hash = "sha256-8ucuEn/oGgbPitL15CDjmo+s1OAVI8bqcwppUWmLp4c=";
+  };
+
+  nativeBuildInputs = [ cmake ];
+
+  meta = with lib; {
+    description = "C++11 multiplatform utility library";
+    homepage = "https://github.com/mosra/corrade";
+    license = with licenses; [
+      mit
+      unlicense
+    ];
+    maintainers = with maintainers; [ SomeoneSerge ];
+    mainProgram = "corrade";
+    platforms = platforms.all;
+  };
+}