1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
|
{
absl-py,
buildPythonPackage,
chex,
dm-env,
dm-tree,
fetchFromGitHub,
flax,
immutabledict,
jax,
lib,
matplotlib,
mediapy,
numpy,
pillow,
pytestCheckHook,
setuptools,
tensorflow,
tqdm,
}:
buildPythonPackage {
pname = "waymax";
version = "0-unstable-2024-03-23";
pyproject = true;
src = fetchFromGitHub {
owner = "waymo-research";
repo = "waymax";
rev = "720f9214a9bf79b3da7926497f0cd0468ca3e630";
hash = "sha256-B1Rp5MATbEelp6G6K2wwV83QpINhOHgvAxb3mBN52Eg=";
};
# AttributeError: jax.tree_map was removed in JAX v0.6.0: use jax.tree.map (jax v0.4.25 or newer) or jax.tree_util.tree_map (any JAX version).
# https://github.com/waymo-research/waymax/pull/77
postPatch = ''
substituteInPlace \
waymax/agents/expert.py \
waymax/agents/waypoint_following_agent.py \
waymax/agents/waypoint_following_agent_test.py \
waymax/dynamics/abstract_dynamics_test.py \
waymax/dynamics/state_dynamics_test.py \
waymax/env/base_environment_test.py \
waymax/env/rollout_test.py \
waymax/env/wrappers/brax_wrapper_test.py \
--replace-fail "jax.tree_map" "jax.tree_util.tree_map"
'';
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 ];
};
}
|
