mmDiff: A Noise-Robust Differentiable Ray-Tracing Framework for mmWave Scene Calibration and Channel Prediction

mmDiff is a differentiable ray-tracing framework designed for mmWave (millimeter-wave) channel prediction and scene calibration. It extends Sionna RT with a noise-robust path-gain model that remains accurate even when the scene geometry is imperfect or noisy — a common challenge when meshes are reconstructed from depth sensors or LiDAR. mmDiff enables gradient-based optimization of radio material parameters directly from channel measurements, bridging physics-based simulation and real-world calibration.

Citation

@misc{mmdiff,
      title={mmDiff: A Noise-Robust Differentiable Ray-Tracing Framework for mmWave Scene Calibration and Channel Prediction}, 
      author={Haofan Lu and Yadi Cao and Wanghao Yi and Omid Abari},
      year={2026},
      eprint={2605.26406},
      archivePrefix={arXiv},
      primaryClass={cs.NI},
      url={https://arxiv.org/abs/2605.26406}, 
}

Installation

Requirements

• Python 3.11
• LLVM for CPU execution via Dr.Jit — see the Dr.Jit LLVM installation guide

Set up the environment

git clone --recurse-submodules https://github.com/LuHaofan/mmdiff.git
cd mmdiff
conda env create -f environment.yml
conda activate mmdiff
pip install -e .

AoA Map Solvers (src/sionna/rt/aoa_map_solvers/)

These modules are developed as part of mmDiff and are not part of the original Sionna RT library. They implement angle-of-arrival (AoA) power spectrum computation via differentiable ray tracing, supporting two antenna paradigms and an optional neural radio material.

File	Description
aoa_map_solver.py	Base ray-tracing solver for synthetic aperture AoA maps. Traces rays from a receiver with a virtual antenna array, accumulating power per steering angle into a 2D (azimuth × elevation) map. Supports LoS, specular reflection, diffuse scattering, and diffraction.
aoa_map.py	Data container and visualizer for synthetic aperture AoA maps. Stores the accumulated path-gain tensor indexed by steering angle and provides show() for polar/Cartesian plots.
mechanical_aoa_map_solver.py	Ray-tracing solver for mechanically-rotated directional antennas. Instead of phase-shifting a virtual array, it physically re-orients a single high-gain antenna pattern over a grid of azimuth/elevation angles and accumulates received power for each orientation.
mechanical_aoa_map.py	Data container and visualizer for mechanical AoA maps. Stores the per-orientation path-gain values and provides show() for 1-D polar scans or 2-D contour plots.
aoa_map_solver_neural.py	Extends AoAMapSolver with a neural radio material. A small MLP (position → permittivity, conductivity, scattering, XPD) is evaluated outside the OptiX kernel and its outputs are passed through SurfaceInteraction fields to the BSDF, keeping the network differentiable for gradient-based material calibration.
mechanical_aoa_map_solver_neural.py	Extends MechanicalAoAMapSolver with the same neural radio material design. Intended for scenarios where a mechanically-scanned directional antenna is used alongside differentiable material optimization.

Radio Materials (src/sionna/rt/radio_materials/)

These modules extend Sionna RT's material system and are developed as part of mmDiff.

File	Description
custom_radio_material.py	Parametric radio material with uniform (spatially-constant) electromagnetic properties. Supports specular reflection, transmission, diffuse scattering, and diffraction. Material parameters (permittivity, conductivity, scattering coefficient, XPD) are directly optimizable via gradient-based methods.
neural_radio_material.py	Neural network-based radio material with spatially-varying properties. Acts as a thin wrapper that stores network architecture configuration; the actual MLP is built and managed by AoAMapSolverNeural, which evaluates it outside the OptiX kernel to avoid memory allocation conflicts inside GPU kernels.

To Do

• Datasets and checkpoints will be released soon

Acknowledgments

mmDiff builds on Sionna RT by NVIDIA (Apache-2.0). The Sionna v0.19.0 baseline is included as a submodule at baselines/sionna (patched fork at github.com/LuHaofan/sionna).

License

Apache-2.0. See LICENSE.

The original Sionna RT source is copyright NVIDIA CORPORATION & AFFILIATES. mmDiff modifications are copyright Haofan Lu.