diff --git a/pkgs/by-name/mu/mujoco/package.nix b/pkgs/by-name/mu/mujoco/package.nix index a5b304f548e4..0803c53ae996 100644 --- a/pkgs/by-name/mu/mujoco/package.nix +++ b/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) diff --git a/pkgs/development/python-modules/mjviser/default.nix b/pkgs/development/python-modules/mjviser/default.nix new file mode 100644 index 000000000000..508a3858097c --- /dev/null +++ b/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 ]; + }; +}) diff --git a/pkgs/development/python-modules/mujoco-warp/default.nix b/pkgs/development/python-modules/mujoco-warp/default.nix new file mode 100644 index 000000000000..aab4d0b83f65 --- /dev/null +++ b/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 ]; + }; +}) diff --git a/pkgs/development/python-modules/rsl-rl-lib/default.nix b/pkgs/development/python-modules/rsl-rl-lib/default.nix new file mode 100644 index 000000000000..debadb023121 --- /dev/null +++ b/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 ]; + }; +}) diff --git a/pkgs/development/python-modules/torchrunx/default.nix b/pkgs/development/python-modules/torchrunx/default.nix new file mode 100644 index 000000000000..42a8aaacbe9f --- /dev/null +++ b/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 ]; + }; +}) diff --git a/pkgs/top-level/python-packages.nix b/pkgs/top-level/python-packages.nix index 35143f8c6dd9..1422960c85c9 100644 --- a/pkgs/top-level/python-packages.nix +++ b/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 { };