Hierarchical NeRF with JAX3D: Volumetric Rendering, Novel-View

Building an End-to-End Hierarchical NeRF with JAX3D

In this tutorial, we build an end-to-end hierarchical Neural Radiance Field (NeRF) using JAX, Flax, Optax, and the volume-rendering primitives provided by JAX3D. As of 2026, NeRF-style approaches remain a foundational technique for neural 3D reconstruction, and JAX3D offers a lightweight, research-grade toolkit for experimenting with them without pulling in heavy framework dependencies.

We begin by constructing a synthetic multi-view dataset from an analytic scene that contains volumetric geometry and view-dependent radiance. Using sample_along_rays and volume_rendering, we establish the forward rendering process that maps rays to pixel colors and depths. We then implement a NeRF architecture featuring positional encoding, skip connections, separate coarse and fine networks, and view-direction conditioning. Hierarchical importance sampling is handled through sample_piecewise_constant_pdf.

Training is performed with JAX JIT compilation, the Adam optimizer, exponential learning-rate decay, and gradient clipping. Finally, we evaluate novel-view synthesis using PSNR, depth and opacity visualization, sampling diagnostics, 360-degree rendering, and marching-cubes geometry extraction.

Setup and Dependencies

The code below installs the required packages and clones the JAX3D repository. To avoid pulling in optional dependencies such as gin and tfds, we load individual modules directly by file path.

import os, sys, subprocess, importlib.util, functools, dataclasses, time, math

def _sh(cmd):
    subprocess.run(cmd, shell=True, check=False,
                   stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)

print("Installing dependencies ...")
_sh(f'{sys.executable} -m pip install -q "etils[array-types,epy,etree,enp]" '
    f'chex flax optax scikit-image')

REPO_DIR = "/content/jax3d" if os.path.isdir("/content") else os.path.abspath("./jax3d")
if not os.path.isdir(REPO_DIR):
    print("Cloning google-research/jax3d ...")
    _sh(f"git clone -q --depth 1 https://github.com/google-research/jax3d.git {REPO_DIR}")

def _load_module_by_path(name, path):
    """Load a single .py file without triggering the parent package __init__.

    `from jax3d.math import volume_rendering` also works if you run
    `pip install .` inside the clone, but that pulls in gin/tfds/etc.
    """
    spec = importlib.util.spec_from_file_location(name, path)
    mod = importlib.util.module_from_spec(spec)
    sys.modules[name] = mod
    spec.loader.exec_module(mod)
    return mod

_VR_PATH = os.path.join(REPO_DIR, "jax3d", ...

What This Tutorial Covers

• Forward rendering: Sampling rays and performing volumetric rendering with JAX3D primitives.
• NeRF architecture: Positional encoding, skip connections, coarse/fine networks, and view-direction conditioning.
• Hierarchical sampling: Importance sampling via sample_piecewise_constant_pdf.
• Training loop: JIT compilation, Adam optimization, learning-rate decay, and gradient clipping.
• Evaluation: PSNR, depth and opacity visualization, sampling diagnostics, 360-degree rendering, and marching-cubes geometry extraction.

By the end, you will have a complete pipeline for volumetric rendering, novel-view synthesis, and 3D reconstruction that runs efficiently on modern accelerators thanks to JAX's just-in-time compilation.

via MarkTechPost