pyrox.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 disabledTests = [
59 # AssertionError: Not equal to tolerance rtol=1e-07, atol=0
60 "test_obs_from_state_for_different_coordinate_system0"
61 "test_obs_from_state_for_different_coordinate_system1"
62 ];
63
64 disabledTestPaths = [
65 # Disable visualization tests that require a GUI
66 # waymax/visualization/viz_test.py Fatal Python error: Aborted
67 "waymax/visualization/viz_test.py"
68 ];
69
70 meta = {
71 description = "JAX-based simulator for autonomous driving research";
72 homepage = "https://github.com/waymo-research/waymax";
73 changelog = "https://github.com/waymo-research/waymax/blob/main/CHANGELOG.md";
74 maintainers = with lib.maintainers; [ samuela ];
75 };
76}