Gradient-Based optimization using Radiomaps

[cyclehome]

Hi all,
in anology to https://nvlabs.github.io/sionna/rt/developer/dev_compat_frameworks.html#training-loop-in-pytorch
I'm trying to set up a primitive optimization that optimizes the tilt of a single transmitter so that the sum of all received-signal-strength values in the observed area is maximized:

import os
os.environ["CUDA_VISIBLE_DEVICES"] = ""

import drjit as dr
import numpy as np
import torch
from sionna.rt import load_scene, PlanarArray, Transmitter, RadioMapSolver


scene = load_scene()
scene.bandwidth = 100e6
scene.frequency = 1815e6
scene.tx_array = PlanarArray(num_rows=8, num_cols=2, pattern="tr38901", polarization="V")

tx_0 = Transmitter(name="tx0", position=[0., 0., 25.], orientation=np.pi/180 * np.array([45., 10., 0.]), power_dbm=23)
scene.add(tx_0)

rm_solver = RadioMapSolver()
rm_solver.loop_mode = "evaluated"  # Needed for gradient computation


@dr.wrap(source="torch", target="drjit")
def compute_map(tilt):
    tx_0.orientation[1] = np.pi/180 * tilt
    rm = rm_solver(scene, max_depth=5, samples_per_tx=10**7, cell_size=(1, 1),
                   size=[400, 400], center=[0, 0, 0], orientation=[0, 0, 0])
    map = rm.transmitter_radio_map("rss", 0)
    rss_sum = dr.sum(map)
    return rss_sum

tilt = torch.tensor(10., requires_grad=True)
optimizer = torch.optim.Adam([tilt], lr=0.05)

num_steps = 3
for step in range(num_steps):
    loss = -compute_map(tilt)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    print("Step: ", step)
    print("map rss sum: ", -loss)
    print("Tilt: ", tilt)
    print("------------------------------------")

The output is as follows:

Step:  0
map rss sum:  tensor(0.0026, grad_fn=<NegBackward0>)
Tilt:  tensor(10.0500, requires_grad=True)
------------------------------------
Step:  1
map rss sum:  tensor(0.0026, grad_fn=<NegBackward0>)
Tilt:  tensor(nan, requires_grad=True)
------------------------------------
Step:  2
map rss sum:  tensor(0., grad_fn=<NegBackward0>)
Tilt:  tensor(nan, requires_grad=True)
------------------------------------
2025-05-09 09:50:43 WARN  main  [Thread] Destructor called while thread 'main' was still running
jit_malloc_shutdown(): leaked
 - host-async memory: 128 B in 2 allocations

So in step 0 the optimizer yields a modified tilt, but already in step 1 it returns NaN. With a different start value, e.g. 0 instead of 10, we get NaN already in step 0.
What am I doing wrong?
Also, what does the "Destructor called while thread 'main' was still running" message tell me? This does not appear with the original example from the documentation (see link above).
Using version 1.0.2.

Thanks a lot!

[faycalaa]

Hi @cyclehome,

Thank you for reporting this issue. We were able to reproduce it, and it will be fixed in the next release.

[faycalaa]

The issue is that dr.acos is not safe to differentiate at -1 and 1, resulting in nan/inf in some cases.

[cyclehome]

Hi @faycalaa,
thanks a lot for the fix, have tested it with the new release. Works much better now.
However, it still fails when using a start value of 0 for the tilt: after the first optimizer.step() the tilt is NaN. Also when starting with a negative value, it works until the tilt needs to cross 0, then it ends up being NaN.
So it seems there is still an issue when the variable is close to 0. Could you have a look at it again?
Thank you!

[faycalaa]

Hi @cyclehome,

Thank you for reporting the issue. It will be fixed in the next release.

[cyclehome]

great, thank you!
Any rough idea when that will be?

[faycalaa]

It seems to come from a call to dr.atan2, which returns NaN when differentiated with both inputs equal to zero.