DOCSEARCHbyAlgolia
Submit Docs

Navix: Accelerated Gridworld Navigation with JAX

repository·main·Indexed 41 days ago

https://github.com/epignatelli/navix

A high-performance JAX-based reimplementation of the MiniGrid environment designed for deep reinforcement learning. Navix enables massive scaling and rapid experimentation by allowing environments to be vectorized with jax.vmap, compiled with jax.jit, and differentiated using jax.grad. It includes a benchmarking system with presets like Navix100K and Navix1M to evaluate algorithms based on training curves, wall-clock time, and computational cost.

Tokens
8.1K
Snippets
19
Records
26
Agent score
51%

What's inside Navix

  1. What is NAVIX?

    main

    NAVIX is a high-throughput reimplementation of the MiniGrid environment suite using JAX. It is designed as a drop-in replacement for the original MiniGrid, optimized for speed and massive parallelization by leveraging JAX's intermediate language representation to run on CPUs, GPUs, or TPUs.

    Key benefits include:

    • High Throughput: Up to 200,000x speedup in batch mode or 20x in unbatched mode.
    • Massive Parallelization: Capable of running thousands of agents (e.g., 2048 PPO agents / 32,768 environments) on a single Nvidia A100 80Gb.
    • JAX Compatibility: Supports jax.vmap, jax.jit, jax.grad, and other JAX transformations.
    • Automatic Differentiation: Enables computing gradients of the environment with respect to the agent's actions.

  2. How Navix benchmarks work

    main

    A Navix benchmark evaluates an algorithm from scratch across a fixed set of environments at a fixed frame budget.

    To participate, you provide an AlgorithmEntry subclass that implements a train method. This method must be compatible with jax.jit/jax.vmap traces (i.e., it cannot return non-computable values like wall-clock time).

    When you call benchmark.run(entry), the system executes the training and returns a BenchmarkResult containing:

    • The TrainingCurve (episodic returns and lengths).
    • wall_time and fps (real measurements using jax.block_until_ready).
    • cost (a CostAnalysis object containing FLOPs, peak memory, and compile time).

    Presets like Navix100K and Navix1M handle the orchestration of training across multiple environments and seeds (defaulting to 16 seeds) using jax.vmap for efficiency.

  3. Core concept: NAVIX environments as PyTreeNodes

    main

    A navix.Environment is implemented as a flax.struct.PyTreeNode. This design choice allows the environment to be fully compatible with the JAX ecosystem, meaning you can apply standard JAX transformations directly to the environment object.

    Specifically, you can use:

    • jax.vmap for vectorizing environments (batching).
    • jax.jit for compiling environment steps into optimized machine code.
    • jax.grad for computing gradients through the environment dynamics.

  4. Use NAVIX as a drop-in replacement for MiniGrid

    main
    NAVIX is designed to be a drop-in replacement for the official MiniGrid environments. You can reuse your existing code and scripts with little to no modification by substituting MiniGrid environment IDs with their corresponding NAVIX IDs.

  5. Install NAVIX

    main

    To use NAVIX, you must first install JAX following the official installation guide for your specific operating system and accelerator (CPU, GPU, or TPU). Once JAX is installed, you can install NAVIX via pip.

    # Install the stable version from PyPI
pip install navix

# Or install the latest version from source
pip install git+https://github.com/epignatelli/navix

  6. Install NAVIX in a conda environment

    main

    It is recommended to install NAVIX within a dedicated conda environment to manage dependencies. Follow these steps to create an environment and install the package:

    conda create -n navix python=3.10
conda activate navix
cd <path/to/navix-repo>
pip install navix

  7. Submit your own benchmark entry

    main

    Follow these steps to submit a new algorithm entry to the Navix benchmarks:

    1. Pick a preset: Choose benchmarks/100K/ or benchmarks/1M/.
    2. Create a folder: Create benchmarks/<preset>/<your-entry>/ (e.g., benchmarks/1M/my-algo/).
    3. Write run.py: Implement the AlgorithmEntry and use the benchmark runner as shown in the implementation guide. The script must be runnable via python run.py with no arguments.
    4. Add config.yml: Create a configuration file in your entry folder with your provenance metadata:
      name: MyAlgo
author: your-github-handle
paper_url: https://arxiv.org/abs/...
navix_commit_url: https://github.com/epignatelli/navix/commit/<sha>
algorithm_commit_url: https://github.com/<you>/<your-repo>/commit/<sha>
      Note: author must be a GitHub handle, and commit URLs must be full URLs ending in a SHA.
    5. Add requirements.txt: Include pinned dependencies required for your run.py to execute.
    6. Run the benchmark: Execute python run.py. This will print a summary and write summary.json, details.json, and diagnostics.npz to your folder.
    7. Submit: Open a Pull Request adding your folder to the repository.

  9. Use the environment interface (step and reset)

    main

    The core interaction loop in NAVIX consists of env.reset(key) and env.step(timestep, action).

    • env.reset(key): Takes a JAX PRNGKey and returns a navix.Timestep object.
    • env.step(timestep, action): Takes the current timestep and an action as arguments. It returns the next timestep object.

    Note that while env.step is stochastic, it does not require a key argument because NAVIX manages the PRNGKey internally. To sample different environment initializations, you must provide different keys to env.reset.

    # Basic interaction loop
timestep = env.reset(key)
for action in actions:
    timestep = env.step(timestep, action)

  10. Run batched environment simulations with vmap

    main

    NAVIX can scale to thousands of parallel environments using jax.vmap. By wrapping a JIT-compiled simulation function (like one using jax.lax.scan) with jax.vmap, you can run multiple simulations in parallel across a batch of keys.

    To implement this:

    1. Generate a batch of keys using jax.random.split(key, num_envs).
    2. Use jax.vmap on your simulation function, specifying in_axes=(0, None) to map over the keys while keeping other arguments (like num_steps) constant.
    3. JIT compile the resulting batched function for maximum performance.
    num_envs = 32
keys = jax.random.split(key, num_envs)

# Create a batched, JIT-compiled simulation function
unroll_batched = jax.jit(
    jax.vmap(unroll_scan, in_axes=(0, None)), 
    static_argnums=(1,)
).lower(keys, 10).compile()

# Run the batch
last_steps = unroll_batched(keys)
print("Batch size of the results", last_steps.reward.shape[0])

  12. Optimize environment execution with JIT compilation

    main

    To achieve high performance, you should use jax.jit to compile your interaction logic. You can compile the individual step function or, for even better performance, compile the entire simulation loop using jax.lax.scan.

    Compiling the entire loop with jax.lax.scan allows XLA to optimize the whole sequence of steps, significantly increasing throughput compared to calling a JIT-compiled step function in a Python loop.

    # Compiling the entire simulation loop with jax.lax.scan
def unroll_scan(key, num_steps=10):
    timestep = env.reset(key)
    actions = jax.random.randint(key, (num_steps,), 0, env.action_space.n)

    timestep, _ = jax.lax.scan(
        lambda timestep, action: (env.step(timestep, action), ()),
        timestep,
        actions,
        unroll=10,
    )
    return timestep

# JIT compile the loop
unroll_jit_loop = jax.jit(unroll_scan, static_argnums=(1,))
timestep = unroll_jit_loop(key, num_steps=10)