Fit a target position by gradient descent

Recover a point target's position from an observed range-azimuth map by differentiating the Radar simulation with PyTorch autograd and Adam.

In this tutorial you solve a small inverse problem. You simulate an "observed" range-azimuth map from a target at a known position, start from a wrong guess, and move the guess with Adam until the simulated map matches. The gradient flows from the loss, through the FFTs, back into the simulated cube and into the target position.

Read Take gradients first if you have not used grad="vjp" yet.

The script#

import torch
from witwin.core import Scene
from witwin.radar import Fmcw, PointTargets, Radar
 
fmcw = Fmcw(slope=60.012e12, sample_rate=4.4e6, samples_per_chirp=256, chirps_per_frame=4,
            adc_start=6e-6, idle=7e-6, ramp_end=65e-6, output="beat")
radar = Radar(carrier=77e9, waveform=fmcw, tx=[(0, 0, 0)], rx=[(i, 0, 0) for i in range(8)],
              antenna_unit="half_wavelength", power=15.0, look_at=(0.0, 0.0, -1.0))
scene = Scene(structures=(), endpoints=[])
 
 
def range_azimuth(positions):
    grad = "vjp" if positions.requires_grad else "none"
    result = radar.simulate(scene, PointTargets(positions=positions, rcs=1.0), times=(0.0,), reflections=0, grad=grad)
    profile = result.frame(0).range_profile(window="hann")    # [TX=1, RX=8, chirp, range]
    beams = torch.fft.fft(profile.data[0], n=64, dim=0)          # azimuth FFT across the 8 receivers
    return beams.abs().mean(dim=1)                               # [azimuth, range] magnitude
 
 
truth = torch.tensor([[0.40, 0.0, -3.00]], device=radar.device)
observed = range_azimuth(truth)
scale = observed.square().sum()
 
guess = torch.tensor([[0.25, 0.0, -2.92]], device=radar.device, requires_grad=True)
optimizer = torch.optim.Adam([guess], lr=0.01)
for step in range(81):
    optimizer.zero_grad()
    loss = (range_azimuth(guess) - observed).square().sum() / scale
    loss.backward()
    optimizer.step()
    if step % 20 == 0:
        print(f"step {step:3d}  loss {loss.item():.4f}  position {[round(v, 3) for v in guess.detach()[0].tolist()]}")
print("truth", truth[0].tolist())

Output on an RTX 5080:

step   0  loss 1.8458  position [0.26, 0.0, -2.93]
step  20  loss 0.0618  position [0.359, 0.0, -3.012]
step  40  loss 0.0252  position [0.394, 0.0, -3.008]
step  60  loss 0.0024  position [0.406, 0.0, -3.002]
step  80  loss 0.0003  position [0.403, 0.0, -3.001]
truth [0.4000000059604645, 0.0, -3.0]

The guess starts 15 cm off in x and 8 cm off in range and ends within 3 mm of the truth.

Design choices#

  1. A linear receive array

    One transmitter and eight receivers half a wavelength apart (antenna_unit="half_wavelength") form a uniform linear array along local x, so an FFT over the receive axis is an azimuth spectrum.

  2. Beat output and a Hann window

    output="beat" returns raw fast-time samples, so range_profile(window="hann") can taper them before the range FFT. The default "spectrum" output is already range-transformed and refuses a fast-time window. The Hann window widens the range peak, which makes the loss smoother.

  3. A magnitude loss

    The carrier phase turns once for every 2 mm of range change at 77 GHz, so a loss on the complex cube has many local minima. The loss here compares magnitudes, which vary smoothly over a range bin (4.3 cm) and an azimuth beam. A magnitude loss gives an exact zero gradient with respect to phase, which is correct.

  4. Gradient only where needed

    grad="vjp" is passed only for the guess. The observed map is computed with grad="none", and the two primals are bitwise identical for the same input.

Tip

A wrong start that is more than a few range bins or beams away gives a loss with no overlap, and the gradient is close to zero. Start within the main lobe, or begin with a coarse grid search.

Things to try#

  • Fit the cross section too: make rcs a 0-dim tensor with requires_grad=True and add it to the optimizer.
  • Add a receiver chain with noise=Noise(figure=10.0) and a fixed seed. The objective stays deterministic, see Take gradients.
  • Replace the azimuth FFT with beam_cube and conventional_steering from witwin.radar.processing, which are differentiable too.