Neural Irradiance Cache (NRC) for Unity URP

A real-time global illumination system based on neural irradiance caching, running on Unity URP (Universal Render Pipeline). It bakes surface irradiance data offline, trains a lightweight neural network to learn the f(position, normal) → irradiance mapping, and performs real-time GPU compute shader inference at runtime, replacing traditional multi-bounce path tracing.

The system also supports an SH Probe Pipeline: densely distributing spherical harmonics probes in 3D space, baking L0-L1 SH coefficients, training a f(position) → SH coefficients(12D) mapping, and reconstructing direction-aware irradiance using surface normals at runtime.

Features

• Three positional encoding schemes: Multi-Resolution Hash Encoding (Instant NGP), Tri-Plane Grid Encoding, Frequency Encoding (NeRF-style)
• Pure PyTorch training with no tiny-cuda-nn dependency (tcnn-accelerated version also available)
• Complete bake workflow inside Unity Editor: Sample → Bake → Export → Train → Load → Inference
• Variable MLP depth (2-4 hidden layers) and width (32/64/128)
• Automatic encoding type detection from weight files, inference shader auto-switching
• Batch training scripts + batch screenshot comparison tools
• SH Probe Pipeline: volumetric spherical harmonics probe baking with L0-L1 direction-aware indirect lighting

Gallery

Test scene: Sponza

Encoding Scheme Comparison

Hash Encoding — Resolution Impact

Hash Encoding — Hidden Dimension Impact

Frequency Encoding — Frequency Band Count Impact

Tri-Plane Encoding — Plane Resolution Impact

SH Probe Mode — Hash Encoding

Dataset: 3,975,860 SH probes (L0+L1, 12D), 8000 steps, batch=65536

Config	Encoding	Encoding Params	MLP	Total Params	File Size	Train Time	Best Loss	Val MAE	Throughput
TP-Small	triplane	64² × 8ch	32×2	100,556	0.19 MB	34.2s	0.002345	0.025445	15.3M/s
TP-Medium	triplane	128² × 16ch	64×2	794,508	1.52 MB	35.0s	0.001284	0.021155	15.0M/s
TP-Large	triplane	256² × 32ch	64×3	6,306,764	12.03 MB	54.7s	0.000667	0.016983	9.6M/s
Hash-Small	hash	8 lvl, T=2^17, res 16→512	32×2	2,099,148	4.00 MB	82.1s	0.001112	0.022266	6.4M/s
Hash-Medium	hash	16 lvl, T=2^19, res 16→2048	64×2	16,784,268	32.01 MB	276.9s	0.000160	0.009382	1.9M/s
Hash-Large	hash	16 lvl, T=2^21, res 16→4096	64×3	67,120,076	128.02 MB	356.9s	0.000019	0.003095	1.5M/s

Small	Medium	Large

SH Probe Mode — Tri-Plane Encoding

Small	Medium	Large

Project Structure

├── Assets/
│   ├── Bake/                          # Bake data and trained weights (.bin)
│   ├── URPRTGI/
│   │   ├── Runtime/
│   │   │   ├── NRCRendererFeature.cs  # URP Renderer Feature, inference entry
│   │   │   ├── NRCWeightLoader.cs     # Weight file loader
│   │   │   ├── RTGIRendererFeature.cs # RTGI path tracer
│   │   │   └── ...
│   │   ├── Editor/
│   │   │   ├── RTGIBakeWindow.cs      # Bake Editor window
│   │   │   ├── SHBakeWindow.cs        # SH probe bake Editor window
│   │   │   ├── SHBakePass.cs          # SH bake dispatch module
│   │   │   ├── SHProbeGenerator.cs    # SH probe spatial generator
│   │   │   ├── NRCBatchScreenshot.cs  # Batch screenshot tool
│   │   │   └── ...
│   │   └── Shaders/
│   │       ├── NRCInference.compute        # Hash encoding inference
│   │       ├── NRCInferenceTriPlane.compute # Tri-plane encoding inference
│   │       ├── NRCInferenceFreqEnc.compute  # Frequency encoding inference
│   │       ├── NRCSHInference.compute       # SH mode inference (12D SH output)
│   │       ├── NRCInferenceMLP.hlsl         # Shared MLP forward pass
│   │       ├── NRCInferenceCommon.hlsl      # Shared utility functions
│   │       ├── SHBake.raytrace             # SH probe full-sphere ray tracing bake
│   │       ├── SHEvaluate.hlsl             # SH basis evaluation (shared)
│   │       ├── SHProbeVisualizer.shader    # SH probe visualization
│   │       └── ...
│   └── Bistro/                        # Test scene (Amazon Lumberyard Bistro)
├── Training/
│   ├── train_nrc.py                   # Main training script (pure PyTorch)
│   ├── train_nrc_tcnn.py             # tiny-cuda-nn accelerated version
│   ├── batch_train_all.ps1           # Windows batch training script
│   ├── batch_train_all.sh            # Linux/macOS batch training script
│   ├── validate_unified.py           # Unified shader validation script
│   └── requirements.txt
├── Screenshots/                       # Batch screenshot output
└── README.md

Quick Start

Requirements

• Unity 2022.3+ (URP) with a DXR-capable GPU
• Python 3.8+, PyTorch >= 2.0, CUDA GPU

1. Bake Training Data

Open Window > RTGI > NRC Bake in Unity:

1. Click Generate Samples to generate surface sample points (default 2M)
2. Click Bake to perform hemisphere integration (path tracing)
3. Click Export to save nrc_bake_data.bin

2. Train

cd Training
pip install -r requirements.txt

# Hash Encoding (recommended, highest quality)
python train_nrc.py -i ../Assets/Bake/nrc_bake_data.bin \
    -o ../Assets/Bake/nrc_weights.bin \
    --encoding hash --steps 5000

# Tri-Plane Grid Encoding (smaller storage)
python train_nrc.py -i ../Assets/Bake/nrc_bake_data.bin \
    -o ../Assets/Bake/nrc_weights_triplane.bin \
    --encoding triplane --plane-res 256 --features-per-plane 16

# Frequency Encoding (lightest, MLP weights only)
python train_nrc.py -i ../Assets/Bake/nrc_bake_data.bin \
    -o ../Assets/Bake/nrc_weights_freq.bin \
    --encoding frequency --n-frequencies 10

3. Load & Inference

Place weight files in Assets/Bake/, select the weight file in the NRC Bake window and click Load. NRCRendererFeature automatically selects the corresponding inference shader based on the encoding_type in the weight file.

SH Probe Bake Workflow

The SH probe pipeline densely distributes spherical harmonics probes in 3D volumetric space. Each probe stores L0-L1 RGB SH coefficients (4×3 = 12 floats), encoding directional irradiance distribution. Unlike the NRC pipeline (surface sampling position+normal → irradiance), the SH pipeline only needs position → SH coefficients, then reconstructs direction-aware indirect lighting using surface normals at runtime.

NRC pipeline:  position + normal → MLP → irradiance(3D)
SH  pipeline:  position → MLP → SH coefficients(12D) → SH evaluate(normal) → irradiance(3D)

1. Generate & Bake SH Probes

Open Window > RTGI > SH Bake in Unity:

1. Configure probe distribution mode (Uniform Grid or Random) and AABB bounds
2. Click Generate SH Probes to generate probe positions in the scene volume
3. Click Bake SH to perform full-sphere ray tracing integration, computing SH coefficients per probe
4. Use the Preview toggle to visualize in Scene View (supports SH Irradiance / Dominant Direction / DC Only modes)
5. Click Export SH Data to save the .bin file

2. Train SH Mode

cd Training

# SH mode training (position → 12D SH coefficients)
python train_nrc.py -i ../Assets/Bake/sh_probe_data.bin \
    -o ../Assets/Bake/nrc_weights_sh.bin \
    --sh_mode --encoding hash --steps 5000

# Also supports triplane and frequency encoding
python train_nrc.py -i ../Assets/Bake/sh_probe_data.bin \
    -o ../Assets/Bake/nrc_weights_sh_triplane.bin \
    --sh_mode --encoding triplane --plane-res 256

3. Runtime Inference

After loading SH weights, NRCRendererFeature auto-detects SH mode (output_dim_pad=16) and dispatches NRCSHInference.compute for inference: position is encoded + MLP outputs 12D SH coefficients, then SHEvaluateIrradiance() reconstructs irradiance using the surface normal.

Encoding Scheme Comparison

Scheme	Encoding Dim	Storage Size	Inference Speed	Quality
Hash Encoding	n_levels × 2 (default 32)	~32MB (hash tables) + MLP	Medium	High
Tri-Plane Encoding	3 × features_per_plane	~few MB (planes) + MLP	Fast	Medium-High
Frequency Encoding	3 × 2 × n_freq	MLP only (~few KB)	Fastest	Medium

Training Parameters

General Parameters

Parameter	Default	Description
--steps	5000	Training steps
--batch-size	65536	Batch size
--hidden-dim	64	MLP hidden layer dimension (32/64/128)
--hidden-layers	2	MLP hidden layer count (2-4)
--lr-encoding	0.01	Encoding parameter learning rate
--lr-network	0.001	MLP learning rate
--jitter-copies	0	Jitter copies per sample for data augmentation
--jitter-radius	0.05	Jitter radius (world space)
--smooth-k	0	KNN spatial smoothing neighbor count (0=disabled)
--smooth-normal-thresh	0.5	KNN smoothing normal similarity threshold

Hash Encoding Parameters

Parameter	Default	Description
--n-levels	16	Number of hash table levels
--log2-hashmap-size	19	Hash table size per level (2^N)
--base-resolution	16	Coarsest resolution
--max-resolution	2048	Finest resolution

Tri-Plane Encoding Parameters

Parameter	Default	Description
--plane-res	128	Plane resolution
--features-per-plane	16	Feature channels per plane

Frequency Encoding Parameters

Parameter	Default	Description
--n-frequencies	10	Number of frequency bands
--include-identity	false	Whether to include raw position

SH Mode Parameters

Parameter	Default	Description
--sh_mode	false	Enable SH probe training mode

In SH mode, the network input is position(3D) and output is 12D SH coefficients (L0-L1 × RGB), with no normal input. The output activation is linear (no exp). Weight export uses output_dim_pad of 16 instead of 4.

Batch Training

The project provides batch training scripts for grid search across parameter combinations of all three encodings:

# Windows (PowerShell)
powershell -ExecutionPolicy Bypass -File Training/batch_train_all.ps1

# Linux/macOS
bash Training/batch_train_all.sh

99 parameter combinations total (72 hash + 18 frequency + 9 triplane), results stored in Assets/Bake/ with parameter-based naming.

Batch Screenshots

Open Window > RTGI > NRC Batch Screenshot in Unity to automatically iterate through all weight files in Assets/Bake/, load each one and capture the Game view, outputting to the Screenshots/ directory.

Weight File Format

56-byte unified header + encoding data + MLP weights (float16):

Offset  Size  Type    Field
0       4     u32     encoding_type (0=hash, 1=triplane, 2=frequency)
4       4     u32     param_a (levels / plane_res / n_frequencies)
8       4     u32     param_b (hash_table_size / 0 / 0)
12      4     u32     param_c (features_per_level / features_per_plane / include_identity)
16      4     f32     res_min
20      4     f32     res_max
24      4     u32     hidden_dim
28      4     u32     hidden_layers
32      24    f32×6   AABB (min_xyz, max_xyz)
56      ...           encoding data (hash tables / tri-plane textures / none)
...     ...           MLP weights (float16, row-major, padded to float4)

Debug Mode

Debug Mode in the NRCRendererFeature Inspector:

Value	Visualization
0	Normal rendering (NRC irradiance)
1	Depth buffer
2	Raw MLP output (pre-sigmoid)
3	Hidden layer activations
4	exp output (no clamp)
5	Normalized position (AABB-mapped RGB)

SH Probe Data File Format

The SH probe export file (version=2) uses a different binary format from NRC bake data (version=1):

Offset  Size        Type        Field
0       4           u32         probe_count
4       12          f32×3       aabb_min (x, y, z)
16      12          f32×3       aabb_max (x, y, z)
28      4           u32         version (= 2)
32      N×12        f32×N×3     positions
32+N×12 N×48        f32×N×12    SH coefficients (per probe: L0_R, L1x_R, L1y_R, L1z_R, L0_G, ..., L0_B, ...)

Total file size = 32 + N × 60 bytes