python3Packages.torchrunx: init at 0.4.0, python3Packages.rsl-rl-lib: init at 5.4.2, python3Packages.mujoco-warp: init at 3.13.0, python3Packages.mjviser: init at 0.0.14, mujoco: install plugins (#563696)

pkgs/by-name/mu/mujoco/package.nix

@@ -108,6 +108,18 @@ stdenv.mkDerivation (finalAttrs: {
     ./mujoco-system-deps-dont-fetch.patch
   ];
 
+  # Install plugins to expose them in python to fix mujoco-warp tests
+  # ref. https://github.com/google-deepmind/mujoco/pull/3602
+  postPatch = ''
+    for plugin in actuator elasticity sensor; do
+      echo "install(TARGETS $plugin)" >> plugin/$plugin/CMakeLists.txt
+    done
+
+    for plugin in sdf usd_decoder; do
+      echo "install(TARGETS $plugin""_plugin)" >> plugin/$plugin/CMakeLists.txt
+    done
+  '';
+
   nativeBuildInputs = [
     cmake
     # git is needed to apply patches to ccd-src and qhull-src (see below)

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

@@ -0,0 +1,64 @@
+{
+  lib,
+  buildPythonPackage,
+  fetchFromGitHub,
+
+  # build-system
+  uv-build,
+
+  # dependencies
+  mujoco,
+  numpy,
+  pillow,
+  trimesh,
+  viser,
+
+  # tests
+  pytestCheckHook,
+}:
+
+buildPythonPackage (finalAttrs: {
+  pname = "mjviser";
+  version = "0.0.14";
+  pyproject = true;
+  __structuredAttrs = true;
+
+  src = fetchFromGitHub {
+    owner = "mujocolab";
+    repo = "mjviser";
+    tag = "v${finalAttrs.version}";
+    hash = "sha256-LzltAgWKK84gQU4IlzKv7SClaUrLCdyU8vzHh4LIqDs=";
+  };
+
+  postPatch = ''
+    substituteInPlace pyproject.toml \
+      --replace-fail \
+        "uv_build>=0.8.19,<0.9.0" \
+        "uv_build"
+  '';
+
+  build-system = [
+    uv-build
+  ];
+
+  dependencies = [
+    mujoco
+    numpy
+    pillow
+    trimesh
+    viser
+  ];
+
+  pythonImportsCheck = [ "mjviser" ];
+
+  nativeCheckInputs = [
+    pytestCheckHook
+  ];
+
+  meta = {
+    description = "Web-based MuJoCo viewer powered by Viser";
+    homepage = "https://github.com/mujocolab/mjviser";
+    license = lib.licenses.asl20;
+    maintainers = with lib.maintainers; [ nim65s ];
+  };
+})

pkgs/development/python-modules/mujoco-warp/default.nix

@@ -0,0 +1,72 @@
+{
+  lib,
+  buildPythonPackage,
+  fetchFromGitHub,
+
+  # build-system
+  setuptools,
+
+  # dependencies
+  absl-py,
+  etils,
+  mujoco,
+  numpy,
+  warp-lang,
+
+  # optional-dependencies
+  jax,
+
+  # tests
+  pytestCheckHook,
+  writableTmpDirAsHomeHook,
+}:
+
+buildPythonPackage (finalAttrs: {
+  pname = "mujoco-warp";
+  version = "3.13.0";
+  pyproject = true;
+  __structuredAttrs = true;
+
+  src = fetchFromGitHub {
+    owner = "google-deepmind";
+    repo = "mujoco_warp";
+    tag = "v${finalAttrs.version}";
+    hash = "sha256-R+YizRSaPOt4eSy6W0dRsFn+ovn+Ks05wVyd5e9lgYc=";
+  };
+
+  build-system = [
+    setuptools
+  ];
+
+  dependencies = [
+    absl-py
+    etils
+    mujoco
+    numpy
+    warp-lang
+  ];
+
+  optional-dependencies = {
+    cpu = [
+      jax
+    ];
+    cuda = [
+      jax
+    ];
+  };
+
+  nativeCheckInputs = [
+    pytestCheckHook
+    writableTmpDirAsHomeHook
+  ];
+
+  pythonImportsCheck = [ "mujoco_warp" ];
+
+  meta = {
+    description = "GPU-optimized version of the MuJoCo physics simulator, designed for NVIDIA hardware";
+    homepage = "https://github.com/google-deepmind/mujoco_warp";
+    changelog = "https://github.com/google-deepmind/mujoco_warp/releases/tag/${finalAttrs.src.tag}";
+    license = lib.licenses.asl20;
+    maintainers = with lib.maintainers; [ nim65s ];
+  };
+})

pkgs/development/python-modules/rsl-rl-lib/default.nix

@@ -0,0 +1,77 @@
+{
+  lib,
+  buildPythonPackage,
+  fetchFromGitHub,
+
+  # build-system
+  setuptools,
+
+  # dependencies
+  gitpython,
+  numpy,
+  onnx,
+  onnxscript,
+  tensorboard,
+  tensordict,
+  torch,
+  torchvision,
+
+  # optional-dependencies
+  wandb,
+
+  # tests
+  pytestCheckHook,
+}:
+
+buildPythonPackage (finalAttrs: {
+  pname = "rsl-rl-lib";
+  version = "5.4.2";
+  pyproject = true;
+  __structuredAttrs = true;
+
+  src = fetchFromGitHub {
+    owner = "leggedrobotics";
+    repo = "rsl_rl";
+    tag = "v${finalAttrs.version}";
+    hash = "sha256-m9M9yWCKs5vLHb7k5A7AFfHko980HDeF3qr1x7s1KGE=";
+  };
+
+  build-system = [
+    setuptools
+  ];
+
+  dependencies = [
+    gitpython
+    numpy
+    onnx
+    onnxscript
+    tensorboard
+    tensordict
+    torch
+    torchvision
+  ];
+
+  optional-dependencies = {
+    # https://github.com/neptune-ai/neptune-client is archived
+    # neptune = [
+    #   neptune
+    # ];
+    wandb = [
+      wandb
+    ];
+  };
+
+  pythonImportsCheck = [ "rsl_rl" ];
+
+  nativeCheckInputs = [
+    pytestCheckHook
+  ];
+
+  meta = {
+    description = "Fast and simple implementation of learning algorithms for robotics";
+    homepage = "https://github.com/leggedrobotics/rsl_rl";
+    changelog = "https://github.com/leggedrobotics/rsl_rl/releases/tag/${finalAttrs.src.tag}";
+    license = lib.licenses.bsd3;
+    maintainers = with lib.maintainers; [ nim65s ];
+  };
+})

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

@@ -0,0 +1,81 @@
+{
+  lib,
+  stdenv,
+  buildPythonPackage,
+  fetchFromGitHub,
+
+  # build-system
+  uv-build,
+
+  # dependencies
+  cloudpickle,
+  fabric,
+  numpy,
+  torch,
+
+  # tests
+  pytestCheckHook,
+  submitit,
+  transformers,
+}:
+
+buildPythonPackage (finalAttrs: {
+  pname = "torchrunx";
+  version = "0.4.0";
+  pyproject = true;
+  __structuredAttrs = true;
+
+  src = fetchFromGitHub {
+    owner = "apoorvkh";
+    repo = "torchrunx";
+    tag = finalAttrs.version;
+    hash = "sha256-cb9X65rnacgB59NdjEVpReFZD1wvC9GWmE0BULnCTow=";
+  };
+
+  build-system = [
+    uv-build
+  ];
+
+  dependencies = [
+    cloudpickle
+    fabric
+    numpy
+    torch
+  ];
+
+  pythonImportsCheck = [ "torchrunx" ];
+
+  nativeCheckInputs = [
+    pytestCheckHook
+    submitit
+    transformers
+  ];
+
+  disabledTests = [
+    # RuntimeError: Not in a SLURM job
+    "test_launch"
+
+    # RuntimeError: Could not detect "srun", are you indeed on a slurm cluster?
+    "test_submitit"
+
+    #  RuntimeError: workers_per_host="gpu", but no GPUs detected on: ['localhost'].
+    "test_distributed_train"
+  ]
+  ++ lib.optionals stdenv.hostPlatform.isDarwin [
+    # torch.distributed.DistNetworkError: The client socket has timed out after 300000ms while
+    # trying to connect to (build03, 63972).
+    "test_error"
+    "test_logging"
+    "test_simple_localhost"
+  ];
+
+  __darwinAllowLocalNetworking = true;
+
+  meta = {
+    description = "Easily run PyTorch on multiple GPUs & machines";
+    homepage = "https://github.com/apoorvkh/torchrunx";
+    changelog = "https://github.com/apoorvkh/torchrunx/releases/tag/${finalAttrs.src.tag}";
+    license = lib.licenses.gpl3Only;
+    maintainers = with lib.maintainers; [ nim65s ];
+  };
+})

pkgs/top-level/python-packages.nix

@@ -10926,6 +10926,8 @@ self: super: with self; {
 
   mizani = callPackage ../development/python-modules/mizani { };
 
+  mjviser = callPackage ../development/python-modules/mjviser { };
+
   mkdocs = callPackage ../development/python-modules/mkdocs { };
 
   mkdocs-autolinks-plugin = callPackage ../development/python-modules/mkdocs-autolinks-plugin { };
@@ -11325,6 +11327,8 @@ self: super: with self; {
 
   mujoco-mjx = callPackage ../development/python-modules/mujoco-mjx { mujoco-main = pkgs.mujoco; };
 
+  mujoco-warp = callPackage ../development/python-modules/mujoco-warp { };
+
   mujson = callPackage ../development/python-modules/mujson { };
 
   mullvad-api = callPackage ../development/python-modules/mullvad-api { };
@@ -18344,6 +18348,8 @@ self: super: with self; {
 
   rsa = callPackage ../development/python-modules/rsa { };
 
+  rsl-rl-lib = callPackage ../development/python-modules/rsl-rl-lib { };
+
   rsskey = callPackage ../development/python-modules/rsskey { };
 
   rst2ansi = callPackage ../development/python-modules/rst2ansi { };
@@ -20962,6 +20968,8 @@ self: super: with self; {
 
   torchrl = callPackage ../development/python-modules/torchrl { };
 
+  torchrunx = callPackage ../development/python-modules/torchrunx { };
+
   torchsde = callPackage ../development/python-modules/torchsde { };
 
   torchsnapshot = callPackage ../development/python-modules/torchsnapshot { };