diff --git a/pkgs/development/python-modules/waymax/default.nix b/pkgs/development/python-modules/waymax/default.nix
new file mode 100644
index 000000000000..1e499b92f03a
--- /dev/null
+++ b/pkgs/development/python-modules/waymax/default.nix
@@ -0,0 +1,70 @@
+{
+  absl-py,
+  buildPythonPackage,
+  chex,
+  dm-env,
+  dm-tree,
+  fetchFromGitHub,
+  flax,
+  immutabledict,
+  jax,
+  lib,
+  matplotlib,
+  mediapy,
+  numpy,
+  pillow,
+  pytestCheckHook,
+  setuptools,
+  tensorflow,
+  tqdm,
+}:
+
+buildPythonPackage rec {
+  pname = "waymax";
+  version = "0-unstable-2025-03-25";
+  pyproject = true;
+
+  src = fetchFromGitHub {
+    owner = "waymo-research";
+    repo = "waymax";
+    rev = "720f9214a9bf79b3da7926497f0cd0468ca3e630";
+    hash = "sha256-B1Rp5MATbEelp6G6K2wwV83QpINhOHgvAxb3mBN52Eg=";
+  };
+
+  build-system = [ setuptools ];
+
+  dependencies = [
+    absl-py
+    chex
+    dm-env
+    dm-tree
+    flax
+    immutabledict
+    jax
+    matplotlib
+    mediapy
+    numpy
+    pillow
+    tensorflow
+    tqdm
+  ];
+
+  nativeCheckInputs = [
+    pytestCheckHook
+  ];
+
+  pythonImportsCheck = [ "waymax" ];
+
+  disabledTestPaths = [
+    # Disable visualization tests that require a GUI
+    # waymax/visualization/viz_test.py Fatal Python error: Aborted
+    "waymax/visualization/viz_test.py"
+  ];
+
+  meta = {
+    description = "A JAX-based simulator for autonomous driving research";
+    homepage = "https://github.com/waymo-research/waymax";
+    changelog = "https://github.com/waymo-research/waymax/blob/main/CHANGELOG.md";
+    maintainers = with lib.maintainers; [ samuela ];
+  };
+}
diff --git a/pkgs/top-level/python-packages.nix b/pkgs/top-level/python-packages.nix
index d42500f8352e..ef1806f770bf 100644
--- a/pkgs/top-level/python-packages.nix
+++ b/pkgs/top-level/python-packages.nix
@@ -18905,6 +18905,8 @@ self: super: with self; {
 
   waybackpy = callPackage ../development/python-modules/waybackpy { };
 
+  waymax = callPackage ../development/python-modules/waymax { };
+
   wazeroutecalculator = callPackage ../development/python-modules/wazeroutecalculator { };
 
   wcag-contrast-ratio = callPackage ../development/python-modules/wcag-contrast-ratio { };