Fit a box to an occupancy grid

This tutorial solves a small inverse problem with Core alone. You have an occupancy grid of an unknown box (here generated, in practice measured or produced by a solver). You start from a guess and let PyTorch move the box until its occupancy matches. It shows how gradients reach geometry parameters, and why the beta of to_mask matters.

It runs in a few seconds on a GPU and in under ten seconds on a recent CPU.

The full script#

import torch
 
from witwin.core import Box
 
device = "cuda" if torch.cuda.is_available() else "cpu"
 
# Sample points: a 64 x 64 x 64 grid over a 2 m cube (spacing about 3.2 cm).
axis = torch.linspace(-1.0, 1.0, 64, device=device)
xx, yy, zz = torch.meshgrid(axis, axis, axis, indexing="ij")
 
# The target: occupancy of a box we pretend to have measured.
target = Box(position=(0.1, -0.2, 0.0), size=(0.8, 0.5, 0.6), device=device)
target_mask = target.to_mask(xx, yy, zz, beta=0.05).detach()
 
# The model: a box whose centre and size are trainable tensors.
center = torch.zeros(3, device=device, requires_grad=True)
size = torch.full((3,), 0.4, device=device, requires_grad=True)
optimizer = torch.optim.Adam([center, size], lr=0.01)
 
for step in range(301):
    model = Box(position=center, size=size)
    mask = model.to_mask(xx, yy, zz, beta=0.05)
    loss = torch.mean((mask - target_mask) ** 2)
    if step % 100 == 0:
        print(
            f"step {step:3d}  loss {loss.item():.2e}  "
            f"center {[round(v, 3) for v in center.tolist()]}  "
            f"size {[round(v, 3) for v in size.tolist()]}"
        )
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

Expected output#

step   0  loss 1.86e-02  center [0.0, 0.0, 0.0]  size [0.4, 0.4, 0.4]
step 100  loss 9.13e-07  center [0.1, -0.199, 0.0]  size [0.802, 0.5, 0.599]
step 200  loss 5.97e-12  center [0.1, -0.2, -0.0]  size [0.8, 0.5, 0.6]
step 300  loss 6.46e-15  center [0.1, -0.2, -0.0]  size [0.8, 0.5, 0.6]

The fit recovers the target centre (0.1, -0.2, 0) and size (0.8, 0.5, 0.6). Numbers in the last digits can differ between CPU and GPU.

How it works#

The model is rebuilt every step. Box(position=center, size=size) is cheap. It stores the two tensors as they are, so the occupancy computed from it depends on center and size in the autograd graph.

to_mask is differentiable. It computes 0.5 * (1 - tanh(signed_distance / beta)). Near the surface the occupancy slides smoothly from 1 to 0, and that slope is what the gradient follows.

beta sets how far the gradient reaches. The gradient is non-zero only within a few beta of the surface. With beta=0.05 (5 cm) the starting box is close enough to every face of the target for each face to feel a pull.

Try the default beta#

Remove both beta=0.05 arguments and run again. to_mask then uses 5 % of the grid spacing, about 1.6 mm, and the fit stalls:

step   0  loss 2.52e-02  center [0.0, 0.0, 0.0]  size [0.4, 0.4, 0.4]
step 100  loss 1.55e-02  center [0.0, -0.19, 0.0]  size [0.455, 0.503, 0.455]
step 200  loss 1.55e-02  center [0.0, -0.19, 0.0]  size [0.456, 0.504, 0.456]
step 300  loss 1.55e-02  center [0.0, -0.19, 0.0]  size [0.457, 0.504, 0.457]

With a thin transition band, a face only gets a gradient where the two masks disagree right at that face. In x and z the starting faces sit inside the target, where both masks are 1, so they never move. In y the +y face starts outside the target, so y gets a gradient and ends close to the target. x and z stay where they started. A common recipe is to start with a wide beta and shrink it as the fit converges.

Where to go from here#

• Fit a Mesh instead: pass a vertex tensor with requires_grad=True and fill_mode="solid" (see Load a mesh). Mesh distances cost more, so use a coarser grid.
• Replace the occupancy loss with a solver output. Because the solvers read the same tensors, the gradient path is the same.