tjh.dev
1{
2 absl-py,
3 buildPythonPackage,
4 chex,
5 dm-env,
6 dm-tree,
7 fetchFromGitHub,
8 flax,
9 immutabledict,
10 jax,
11 lib,
12 matplotlib,
13 mediapy,
14 numpy,
15 pillow,
16 pytestCheckHook,
17 setuptools,
18 tensorflow,
19 tqdm,
20}:
21
22buildPythonPackage {
23 pname = "waymax";
24 version = "0-unstable-2025-05-30";
25 pyproject = true;
26
27 src = fetchFromGitHub {
28 owner = "waymo-research";
29 repo = "waymax";
30 rev = "48b33d71aac20a22db7d25f2d3220596899d944a";
31 hash = "sha256-YV0KI0UrFXO3HvKjqJE+K+hJJuYI4GiIR4l1fZNnl/E=";
32 };
33
34 build-system = [ setuptools ];
35
36 dependencies = [
37 absl-py
38 chex
39 dm-env
40 dm-tree
41 flax
42 immutabledict
43 jax
44 matplotlib
45 mediapy
46 numpy
47 pillow
48 tensorflow
49 tqdm
50 ];
51
52 nativeCheckInputs = [
53 pytestCheckHook
54 ];
55
56 pythonImportsCheck = [ "waymax" ];
57
58 disabledTestPaths = [
59 # Disable visualization tests that require a GUI
60 # waymax/visualization/viz_test.py Fatal Python error: Aborted
61 "waymax/visualization/viz_test.py"
62 ];
63
64 meta = {
65 description = "JAX-based simulator for autonomous driving research";
66 homepage = "https://github.com/waymo-research/waymax";
67 changelog = "https://github.com/waymo-research/waymax/blob/main/CHANGELOG.md";
68 maintainers = with lib.maintainers; [ samuela ];
69 };
70}