F1TENTH Gym JAX

This repository contains a JAX-compatible multi-agent F1TENTH racing environment. The main API is f1tenth_gym_jax.make(...), which returns a jittable environment with reset(key) and step(key, state, actions) methods.

The project is under active development.

Quickstart

Install the package in an isolated Python 3.11-3.13 environment with uv:

git clone https://github.com/f1tenth/f1tenth_gym_jax.git
cd f1tenth_gym_jax
uv sync

uv is the official install path for this repository because it resolves the git-backed jax-pf dependency and applies the dependency overrides recorded in pyproject.toml and uv.lock.

Optional extras are split by workflow:

uv sync --extra examples  # plotting, web dashboard, and track generation examples
uv sync --extra rl        # PPO training/evaluation dependencies
uv sync --extra docs      # Sphinx documentation build
uv sync --extra cuda      # JAX CUDA 13 support

The default install uses CPU JAX wheels. The cuda extra follows JAX's jax[cuda13] packaging and requires a compatible Linux NVIDIA driver.

Run a minimal rollout:

import jax
import jax.numpy as jnp

from f1tenth_gym_jax import make

env = make("Spielberg_1_noscan_nocollision_progress_acceleration+steeringvelocity_1_100_v0")
key = jax.random.key(0)
obs, state = env.reset(key)

actions = {"agent_0": jnp.array([0.0, 1.0])}  # [steering_velocity, acceleration]
key, step_key = jax.random.split(key)
obs, state, rewards, dones, infos = env.step(step_key, state, actions)

Longer example usage lives in:

• examples/train_ppo_example.py
• examples/eval_ppo_example.py
• examples/waypoint_follow.py
• examples/mppi_example.py
• examples/render_dashboard.py
• examples/run_in_empty_track.py
• examples/random_trackgen.py
• examples/benchmark_example.ipynb
• examples/rendering_example.ipynb

Visualization

Rollout examples can write a self-contained HTML dashboard:

uv run python examples/waypoint_follow.py \
  --num-agents 3 \
  --num-envs 10 \
  --steps 500 \
  --render-output /tmp/f1tenth_dashboard.html

If the browser is running on the same machine, open the generated HTML file directly. If the browser is on another machine, serve the output directory from the remote host:

uv run python -m http.server 8766 --bind 0.0.0.0 --directory /tmp

Then open:

http://remote-host:8766/f1tenth_dashboard.html

Environment IDs

Environment IDs use this pattern:

{map}_{num_agents}_{scan|noscan}_{collision|nocollision}_{rewards}_{longitudinal+steering}_{timestep_ratio}_{max_steps}_v0

For the default episode length, use the shorthand form without the max_steps field:

{map}_{num_agents}_{scan|noscan}_{collision|nocollision}_{rewards}_{longitudinal+steering}_{timestep_ratio}_v0

Examples:

Spielberg_1_scan_collision_progress+alive_velocity+steeringangle_10_v0
Spielberg_1_scan_collision_progress+alive_velocity+steeringangle_10_500_v0

The final v0 is the environment ID version. Older model filenames that omit timestep_ratio are still accepted and use timestep_ratio=1.

Action vectors are ordered as [steering_command, longitudinal_command]. For example, acceleration+steeringvelocity environments expect [steering_velocity, acceleration].

Bundled maps are loaded from the installed package. Downloaded maps are cached under $XDG_CACHE_HOME/f1tenth_gym_jax/maps by default; set F1TENTH_GYM_JAX_MAP_DIR to use another writable map directory.

Docker

The default Docker image installs the standard dependency set. Rollout visualization is generated as a standalone HTML dashboard that can be opened in any browser.

docker build -t f1tenth_gym_jax -f Dockerfile .
docker run -it f1tenth_gym_jax

Citing

If you find this environment useful, please consider citing:

@inproceedings{okelly2020f1tenth,
  title={F1TENTH: An Open-source Evaluation Environment for Continuous Control and Reinforcement Learning},
  author={O'Kelly, Matthew and Zheng, Hongrui and Karthik, Dhruv and Mangharam, Rahul},
  booktitle={NeurIPS 2019 Competition and Demonstration Track},
  pages={77--89},
  year={2020},
  organization={PMLR}
}