JaxNaviSim: A GPU-Accelerated, Data-Driven Simulator and Learning Benchmark for Maritime Navigation

Coordinating vessels through the shared waterways of busy ports grows more challenging as vessel traffic increases. Port authorities need realistic simulations to assess how arrival schedules, anchorage assignments, and routes affect vessel throughput and navigation safety. We introduce JaxNaviSim, a data-driven simulator and learning benchmark coupling Automatic Identification System (AIS) data with fast, controllable traffic simulation. First, we release a dataset of  141 million AIS records from over 30K vessels across eight high-density regions of the Singapore Strait, among the world's busiest waterways, with  6.7 million vessel-hours. Unlike several existing AIS datasets which mostly cover open seas, ours captures close-quarters vessel traffic in the narrow strait. Second, we formulate navigation as a goal-reaching task, and implement a scalable JAX simulator supporting historical replay, selective vessel control with batched GPU execution. We also incorporate multiple vessel dynamics models and sea surface currents. Finally, we benchmark imitation learning (IL), offline and online RL, with ablations of key design choices. Learned policies reach high goal completion, but collisions persist. We therefore propose Predictive Action Refinement (PAR), which refines actions using short-horizon interaction forecasts, substantially reducing collisions. Our contributions enable research on scarce maritime resource allocation, traffic modeling, digital twins, and provide large real-world testbed for IL/RL methods.