diff --git a/benchmarks/play_model_benchmark.py b/benchmarks/complex_rail_benchmark.py similarity index 91% rename from benchmarks/play_model_benchmark.py rename to benchmarks/complex_rail_benchmark.py index 41c15ef0ef1cbbeebcf547ffa967dc9c31fecbe6..1f6985dcdb374c09f5857e44c4fc484bd62d9c6e 100644 --- a/benchmarks/play_model_benchmark.py +++ b/benchmarks/complex_rail_benchmark.py @@ -1,3 +1,4 @@ +"""Run benchmarks on complex rail flatland.""" import random import numpy as np @@ -7,7 +8,8 @@ from flatland.envs.generators import complex_rail_generator from flatland.envs.rail_env import RailEnv -def main(): +def run_benchmark(): + """Run benchmark on a small number of agents in complex rail environment.""" random.seed(1) np.random.seed(1) @@ -66,4 +68,4 @@ if __name__ == "__main__": with Benchmarker(cycle=20, extra=1) as bench: @bench("Everything") def _(bm): - main() + run_benchmark()