Hierarchical NeRF with JAX3D for Volumetric Rendering, Novel-View Synthesis, and 3D Reconstruction

Hierarchical NeRF with JAX3D for Volumetric Rendering, Novel-View Synthesis, and 3D Reconstruction

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. We first construct a synthetic multi-view dataset from an analytic scene containing volumetric geometry and view-dependent radiance, using sample_along_rays and volume_rendering to establish the forward rendering process. We then implement a NeRF with positional encoding, skip connections, separate coarse and fine networks, and view-direction conditioning, followed by hierarchical importance sampling through sample_piecewise_constant_pdf. We train the model with JAX JIT compilation, Adam optimization, exponential learning-rate decay, and gradient clipping, and finally evaluate novel-view synthesis using PSNR, depth and opacity visualization, sampling diagnostics, 360-degree rendering, and marching-cubes geometry extraction.

Copy CodeCopiedUse a different Browser

We set up the JAX3D environment, install the required dependencies, and load the volume_rendering module directly from the cloned repository. We configure GPU/CPU-adaptive training parameters and establish the camera model using pinhole intrinsics, look-at poses, and orbit-based camera placement. We then generate normalized world-space rays from each camera pose, providing the geometric foundation for the rendering pipeline.

We construct an analytic ground-truth scene containing soft-edged spheres, a patterned floor, and view-dependent specular radiance. We render this scene with JAX3Ds volume-rendering implementation to generate consistent RGB observations, depths, and opacity values across multiple camera views. We organize the resulting images into a flattened ray pool so that we can efficiently sample random rays during NeRF training.

We implement the NeRF representation using sinusoidal positional encoding for both spatial coordinates and viewing directions. We use a deep Flax MLP with a skip connection to predict non-negative volumetric density from position while conditioning RGB on the viewing direction. We therefore separate view-independent geometry from view-dependent appearance, allowing the model to represent both scene structure and specular effects.

We implement the core hierarchical renderer by first sampling coarse points along each ray and compositing their densities and colors through JAX3Ds volume-rendering operator. We convert the resulting coarse rendering weights into a piecewise-constant probability distribution and importance-sample additional fine points around high-contribution regions. We combine and sort the coarse and fine samples before performing the final fine-network rendering, while stopping gradients through the sampling operation.

We initialize independent coarse and fine NeRF networks and optimize them jointly with Adam, using exponential learning-rate decay and global gradient clipping. We supervise both rendering stages against ground-truth ray colors, encouraging the coarse network to learn useful sampling distributions while improving the final fine reconstruction. We run the training step with JAX JIT compilation and monitor the fine-network PSNR throughout optimization.

We evaluate the trained representation through chunked novel-view rendering and measure reconstruction quality with held-out PSNR, along with depth and opacity maps. We visualize how hierarchical sampling concentrates fine samples around important surfaces, then generate a 360-degree orbit GIF to inspect the learned radiance field from multiple viewpoints. We finally query the learned density on a 3D grid and apply marching cubes to extract an approximate geometric isosurface.

In conclusion, we demonstrated the complete inverse-rendering pipeline by learning a continuous density and radiance field from synthetic multi-view observations and reconstructing it through hierarchical volume rendering. We used the coarse network to identify informative regions along each ray and the fine network to concentrate additional samples around high-contribution surfaces. At the same time, view-direction encoding allows us to model view-dependent appearance. In the final evaluation stages, we measured novel-view reconstruction quality with PSNR, inspected learned depth and opacity, visualized importance-sampling behavior, generated a 360-degree orbit, and extracted an approximate learned geometry with marching cubes. Overall, we showed how the mathematical components of jax3d integrate with modern JAX-based neural-network training to form a compact yet technically complete NeRF reconstruction system.

Check out the FULL CODES here. All credit goes to the researcher of this project. Also, feel free to follow us on Twitter and don’t forget to join our 150k+ML SubReddit and Subscribe to our Newsletter. Wait! are you on telegram? now you can join us on telegram as well.

Need to partner with us for promoting your GitHub Repo OR Hugging Face Page OR Product Release OR Webinar etc.? Connect with us

Previous articleA Princeton Researcher Proposes Recurrent Looped Transformer (RLT) that Carries Decoder State across Every Token, Fixing 96 Blocks per Token with Unbounded Temporal Depth

Next articleAnthropics 3-Step Pace the Frontier Plan Wins OpenAI, xAI and Microsoft Support: Is It Too Late to Slow AI Down?