Assembly Optimization Strategies

println("Julia ", VERSION)

Julia 1.12.6

Description

The built-in assemble drives the poses with LBFGS over a fixed sample of the reference surface. For a modest shuffle that is enough, but a scattered shuffle can trap it in a local minimum: the surface RMS is blind to the interior, and the fixed sample gives no signal in regions the pieces have vacated.

This tutorial explores two more robust ways to drive the same pose optimization — both replace the surface sample with points drawn inside the bounding box of the reference and the pieces, and both are shown against the sliced-box problem:

Recenter → solve → resample — undo the bulk displacement by centering the pieces on the reference, solve on a volumetric point cloud, then shrink the box and resample, repeating a few times (coarse-to-fine).
Stochastic Adam — forget the batch solver; take small Adam steps on a fresh random point cloud every iteration, re-deriving the bounding box each time.

The code blocks below are written to be copied and run; they are not executed here, so printed values are illustrative.

using Rasmah
using LinearAlgebra
using Random: Random

Setup: the sliced box

Same problem as the self-assembly tutorial, but shuffled harder so the plain LBFGS approach would struggle:

Random.seed!(42)

B = to_sdf(box(2.0, 2.0, 2.0))
slab(xc, hx) = intersection(B, translate(to_sdf(box(2hx, 2.0, 2.0)), xc, 0.0, 0.0))
base = [slab(-0.75, 0.25), slab(-0.25, 0.25), slab(0.25, 0.25), slab(0.75, 0.25)]
np = length(base)

shuffles = [
    ([0.60, -0.45, 0.35], rotation_z(0.40)),
    ([-0.55, 0.50, -0.30], rotation_exp(skew([0.30, 0.20, 0.10]))),
    ([0.35, 0.55, 0.50], rotation_y(0.35)),
    ([-0.50, -0.40, -0.45], rotation_x(0.30)),
]
pieces = [rotate(translate(p, t...), R) for (p, (t, R)) in zip(base, shuffles)]

4-element Vector{RotateSDF{TranslateSDF{IntersectionSDF{BoxSDF{Float64, Float64, Float64}, TranslateSDF{BoxSDF{Float64, Float64, Float64}, Vector{Float64}}}, Vector{Float64}}, Matrix{Float64}}}:
 RotateSDF{TranslateSDF{IntersectionSDF{BoxSDF{Float64, Float64, Float64}, TranslateSDF{BoxSDF{Float64, Float64, Float64}, Vector{Float64}}}, Vector{Float64}}, Matrix{Float64}}(TranslateSDF{IntersectionSDF{BoxSDF{Float64, Float64, Float64}, TranslateSDF{BoxSDF{Float64, Float64, Float64}, Vector{Float64}}}, Vector{Float64}}(IntersectionSDF{BoxSDF{Float64, Float64, Float64}, TranslateSDF{BoxSDF{Float64, Float64, Float64}, Vector{Float64}}}(BoxSDF{Float64, Float64, Float64}(1.0, 1.0, 1.0), TranslateSDF{BoxSDF{Float64, Float64, Float64}, Vector{Float64}}(BoxSDF{Float64, Float64, Float64}(0.25, 1.0, 1.0), [-0.75, 0.0, 0.0])), [0.6, -0.45, 0.35]), [0.9210609940028851 -0.3894183423086505 0.0; 0.3894183423086505 0.9210609940028851 0.0; 0.0 0.0 1.0])
 RotateSDF{TranslateSDF{IntersectionSDF{BoxSDF{Float64, Float64, Float64}, TranslateSDF{BoxSDF{Float64, Float64, Float64}, Vector{Float64}}}, Vector{Float64}}, Matrix{Float64}}(TranslateSDF{IntersectionSDF{BoxSDF{Float64, Float64, Float64}, TranslateSDF{BoxSDF{Float64, Float64, Float64}, Vector{Float64}}}, Vector{Float64}}(IntersectionSDF{BoxSDF{Float64, Float64, Float64}, TranslateSDF{BoxSDF{Float64, Float64, Float64}, Vector{Float64}}}(BoxSDF{Float64, Float64, Float64}(1.0, 1.0, 1.0), TranslateSDF{BoxSDF{Float64, Float64, Float64}, Vector{Float64}}(BoxSDF{Float64, Float64, Float64}(0.25, 1.0, 1.0), [-0.25, 0.0, 0.0])), [-0.55, 0.5, -0.3]), [0.9752903089530457 -0.06803131640494002 0.21019170595074288; 0.12733457491763028 0.9505806179060915 -0.28316496056507373; -0.18054007669439776 0.3029327134026371 0.9357548032779188])
 RotateSDF{TranslateSDF{IntersectionSDF{BoxSDF{Float64, Float64, Float64}, TranslateSDF{BoxSDF{Float64, Float64, Float64}, Vector{Float64}}}, Vector{Float64}}, Matrix{Float64}}(TranslateSDF{IntersectionSDF{BoxSDF{Float64, Float64, Float64}, TranslateSDF{BoxSDF{Float64, Float64, Float64}, Vector{Float64}}}, Vector{Float64}}(IntersectionSDF{BoxSDF{Float64, Float64, Float64}, TranslateSDF{BoxSDF{Float64, Float64, Float64}, Vector{Float64}}}(BoxSDF{Float64, Float64, Float64}(1.0, 1.0, 1.0), TranslateSDF{BoxSDF{Float64, Float64, Float64}, Vector{Float64}}(BoxSDF{Float64, Float64, Float64}(0.25, 1.0, 1.0), [0.25, 0.0, 0.0])), [0.35, 0.55, 0.5]), [0.9393727128473789 0.0 0.34289780745545134; 0.0 1.0 0.0; -0.34289780745545134 0.0 0.9393727128473789])
 RotateSDF{TranslateSDF{IntersectionSDF{BoxSDF{Float64, Float64, Float64}, TranslateSDF{BoxSDF{Float64, Float64, Float64}, Vector{Float64}}}, Vector{Float64}}, Matrix{Float64}}(TranslateSDF{IntersectionSDF{BoxSDF{Float64, Float64, Float64}, TranslateSDF{BoxSDF{Float64, Float64, Float64}, Vector{Float64}}}, Vector{Float64}}(IntersectionSDF{BoxSDF{Float64, Float64, Float64}, TranslateSDF{BoxSDF{Float64, Float64, Float64}, Vector{Float64}}}(BoxSDF{Float64, Float64, Float64}(1.0, 1.0, 1.0), TranslateSDF{BoxSDF{Float64, Float64, Float64}, Vector{Float64}}(BoxSDF{Float64, Float64, Float64}(0.25, 1.0, 1.0), [0.75, 0.0, 0.0])), [-0.5, -0.4, -0.45]), [1.0 0.0 0.0; 0.0 0.955336489125606 -0.29552020666133955; 0.0 0.29552020666133955 0.955336489125606])

The shuffled pieces:

vtk(reduce(union, pieces); n=64)

Zoom Opacity

1. A volumetric (occupancy) loss

Instead of the surface gap, compare the two solids point by point. With a smooth Heaviside H(-f) = 1/(1 + e^{βf}) as an occupancy indicator, the squared occupancy difference

\[L(\theta) = \frac{1}{N}\sum_{x \in \text{box}} \big(H(-A_\theta(x)) - H(-B(x))\big)^2\]

is zero exactly when the union of the moved pieces A_θ coincides with the box B as a set — so internal gaps, overlaps, and protruding pieces all raise it. It is a sum of per-point terms, which is exactly what both strategies below need.

H(s; β = 20.0) = 1 / (1 + exp(-β * s))
occupancy(g, x; β = 20.0) = H(-g(x); β)

move(pieces, θ) = [
    rotate(translate(pieces[i], θ[6i - 5], θ[6i - 4], θ[6i - 3]),
           rotation_exp(skew(θ[(6i - 2):6i]))) for i in eachindex(pieces)
]
assemble_sdf(pieces, θ) = SoftUnion(move(pieces, θ), 0.05)

function bbox(pieces, θ)
    lo1, hi1 = bounds(B)
    lo2, hi2 = bounds(assemble_sdf(pieces, θ))
    min.(lo1, lo2), max.(hi1, hi2)
end

sample_points(lo, hi, n) = [rand(3) .* (hi .- lo) .+ lo for _ in 1:n]

function occupancy_loss(pieces, pts, θ; β = 20.0)
    A = assemble_sdf(pieces, θ)
    s = zero(eltype(θ))
    for x in pts
        d = occupancy(A, x; β) - occupancy(B, x; β)
        s += d * d
    end
    s / length(pts)
end

occupancy_loss (generic function with 1 method)

The reference points x are ordinary Float64 vectors; only the pose θ is differentiated, so AD flows through move/assemble_sdf/occupancy and back into the optimizer.

2. Strategy 1 — recenter, solve, resample (coarse-to-fine)

A shuffled piece has no idea where it belongs, and LBFGS has no idea how far to travel. The first trick is to remove the bulk displacement up front: move the pieces so their combined centroid sits on the reference centroid.

function centroid(g, lo, hi; n = 24, β = 20.0)
    step = (hi .- lo) ./ n
    num = zeros(3)
    den = 0.0
    for k in 0:(n - 1), j in 0:(n - 1), i in 0:(n - 1)
        x = lo .+ [i + 0.5, j + 0.5, k + 0.5] .* step
        w = occupancy(g, x; β)
        num .+= w .* x
        den += w
    end
    den > 0 ? num ./ den : (lo .+ hi) ./ 2
end

lo0, hi0 = bbox(pieces, zeros(6np))
c_ref = centroid(B, lo0, hi0)
c_parts = sum(centroid(p, bounds(p)...) for p in pieces) / np
offset = c_ref - c_parts
pieces = [translate(p, offset...) for p in pieces]

4-element Vector{TranslateSDF{RotateSDF{TranslateSDF{IntersectionSDF{BoxSDF{Float64, Float64, Float64}, TranslateSDF{BoxSDF{Float64, Float64, Float64}, Vector{Float64}}}, Vector{Float64}}, Matrix{Float64}}, Vector{Float64}}}:
 TranslateSDF{RotateSDF{TranslateSDF{IntersectionSDF{BoxSDF{Float64, Float64, Float64}, TranslateSDF{BoxSDF{Float64, Float64, Float64}, Vector{Float64}}}, Vector{Float64}}, Matrix{Float64}}, Vector{Float64}}(RotateSDF{TranslateSDF{IntersectionSDF{BoxSDF{Float64, Float64, Float64}, TranslateSDF{BoxSDF{Float64, Float64, Float64}, Vector{Float64}}}, Vector{Float64}}, Matrix{Float64}}(TranslateSDF{IntersectionSDF{BoxSDF{Float64, Float64, Float64}, TranslateSDF{BoxSDF{Float64, Float64, Float64}, Vector{Float64}}}, Vector{Float64}}(IntersectionSDF{BoxSDF{Float64, Float64, Float64}, TranslateSDF{BoxSDF{Float64, Float64, Float64}, Vector{Float64}}}(BoxSDF{Float64, Float64, Float64}(1.0, 1.0, 1.0), TranslateSDF{BoxSDF{Float64, Float64, Float64}, Vector{Float64}}(BoxSDF{Float64, Float64, Float64}(0.25, 1.0, 1.0), [-0.75, 0.0, 0.0])), [0.6, -0.45, 0.35]), [0.9210609940028851 -0.3894183423086505 0.0; 0.3894183423086505 0.9210609940028851 0.0; 0.0 0.0 1.0]), [-0.036526766731538915, -0.07155314805286443, -0.019736968746098284])
 TranslateSDF{RotateSDF{TranslateSDF{IntersectionSDF{BoxSDF{Float64, Float64, Float64}, TranslateSDF{BoxSDF{Float64, Float64, Float64}, Vector{Float64}}}, Vector{Float64}}, Matrix{Float64}}, Vector{Float64}}(RotateSDF{TranslateSDF{IntersectionSDF{BoxSDF{Float64, Float64, Float64}, TranslateSDF{BoxSDF{Float64, Float64, Float64}, Vector{Float64}}}, Vector{Float64}}, Matrix{Float64}}(TranslateSDF{IntersectionSDF{BoxSDF{Float64, Float64, Float64}, TranslateSDF{BoxSDF{Float64, Float64, Float64}, Vector{Float64}}}, Vector{Float64}}(IntersectionSDF{BoxSDF{Float64, Float64, Float64}, TranslateSDF{BoxSDF{Float64, Float64, Float64}, Vector{Float64}}}(BoxSDF{Float64, Float64, Float64}(1.0, 1.0, 1.0), TranslateSDF{BoxSDF{Float64, Float64, Float64}, Vector{Float64}}(BoxSDF{Float64, Float64, Float64}(0.25, 1.0, 1.0), [-0.25, 0.0, 0.0])), [-0.55, 0.5, -0.3]), [0.9752903089530457 -0.06803131640494002 0.21019170595074288; 0.12733457491763028 0.9505806179060915 -0.28316496056507373; -0.18054007669439776 0.3029327134026371 0.9357548032779188]), [-0.036526766731538915, -0.07155314805286443, -0.019736968746098284])
 TranslateSDF{RotateSDF{TranslateSDF{IntersectionSDF{BoxSDF{Float64, Float64, Float64}, TranslateSDF{BoxSDF{Float64, Float64, Float64}, Vector{Float64}}}, Vector{Float64}}, Matrix{Float64}}, Vector{Float64}}(RotateSDF{TranslateSDF{IntersectionSDF{BoxSDF{Float64, Float64, Float64}, TranslateSDF{BoxSDF{Float64, Float64, Float64}, Vector{Float64}}}, Vector{Float64}}, Matrix{Float64}}(TranslateSDF{IntersectionSDF{BoxSDF{Float64, Float64, Float64}, TranslateSDF{BoxSDF{Float64, Float64, Float64}, Vector{Float64}}}, Vector{Float64}}(IntersectionSDF{BoxSDF{Float64, Float64, Float64}, TranslateSDF{BoxSDF{Float64, Float64, Float64}, Vector{Float64}}}(BoxSDF{Float64, Float64, Float64}(1.0, 1.0, 1.0), TranslateSDF{BoxSDF{Float64, Float64, Float64}, Vector{Float64}}(BoxSDF{Float64, Float64, Float64}(0.25, 1.0, 1.0), [0.25, 0.0, 0.0])), [0.35, 0.55, 0.5]), [0.9393727128473789 0.0 0.34289780745545134; 0.0 1.0 0.0; -0.34289780745545134 0.0 0.9393727128473789]), [-0.036526766731538915, -0.07155314805286443, -0.019736968746098284])
 TranslateSDF{RotateSDF{TranslateSDF{IntersectionSDF{BoxSDF{Float64, Float64, Float64}, TranslateSDF{BoxSDF{Float64, Float64, Float64}, Vector{Float64}}}, Vector{Float64}}, Matrix{Float64}}, Vector{Float64}}(RotateSDF{TranslateSDF{IntersectionSDF{BoxSDF{Float64, Float64, Float64}, TranslateSDF{BoxSDF{Float64, Float64, Float64}, Vector{Float64}}}, Vector{Float64}}, Matrix{Float64}}(TranslateSDF{IntersectionSDF{BoxSDF{Float64, Float64, Float64}, TranslateSDF{BoxSDF{Float64, Float64, Float64}, Vector{Float64}}}, Vector{Float64}}(IntersectionSDF{BoxSDF{Float64, Float64, Float64}, TranslateSDF{BoxSDF{Float64, Float64, Float64}, Vector{Float64}}}(BoxSDF{Float64, Float64, Float64}(1.0, 1.0, 1.0), TranslateSDF{BoxSDF{Float64, Float64, Float64}, Vector{Float64}}(BoxSDF{Float64, Float64, Float64}(0.25, 1.0, 1.0), [0.75, 0.0, 0.0])), [-0.5, -0.4, -0.45]), [1.0 0.0 0.0; 0.0 0.955336489125606 -0.29552020666133955; 0.0 0.29552020666133955 0.955336489125606]), [-0.036526766731538915, -0.07155314805286443, -0.019736968746098284])

After recentering the pieces still point every which way, but their centre of mass is right. Then loop: sample the current bounding box, solve to convergence, recompute the (tighter) box, resample — without recentering again — and repeat. Each round the box shrinks toward the reference, so the same number of points resolves the boundary more finely:

function solve_round(pieces, θ, pts; iterations = 60)
    f(θ) = occupancy_loss(pieces, pts, θ)
    g = θ -> gradient(f, θ)
    res = minimize(ObjectiveFunction(f, g), collect(Float64, θ); iterations = iterations)
    res.minimizer
end

θ = zeros(6np)
for r in 1:4
    lo, hi = bbox(pieces, θ)
    pts = sample_points(lo, hi, 4000)
    θ = solve_round(pieces, θ, pts)
    println("round $r: loss = ", round(occupancy_loss(pieces, pts, θ); digits = 6))
end

round 1: loss = 0.003771
round 2: loss = 0.002925
round 3: loss = 0.003178
round 4: loss = 0.002888

The recovered box:

assembled = reduce(union, [
    rotate(translate(pieces[i], θ[6i - 5], θ[6i - 4], θ[6i - 3]), rotation_exp(skew(θ[(6i - 2):6i])))
    for i in eachindex(pieces)
])
println("reassembled volume = ", round(volume(surface_mesh(assembled; n=32)); digits=3), "  (want 8)")
vtk(assembled; n=64)

reassembled volume = 7.655 m^3  (want 8)

Zoom Opacity

3. Strategy 2 — stochastic Adam with per-step resampling

Batch LBFGS is memory-hungry and greedy. An alternative is Adam: keep first/second moment buffers for the 6n pose parameters and take a small step per iteration, drawing a fresh random point cloud from the (recomputed) bounding box every step. The resampling acts as both a stochastic gradient estimate and a coarse-to-fine schedule in one.

θ = zeros(6np)
m = zeros(6np)
v = zeros(6np)
α, β1, β2, ϵ = 0.01, 0.9, 0.999, 1e-8

for t in 1:300
    lo, hi = bbox(pieces, θ)
    pts = sample_points(lo, hi, 256)          # fresh mini-batch every step
    g = gradient(θ -> occupancy_loss(pieces, pts, θ), θ)

    m .= β1 .* m .+ (1 - β1) .* g
    v .= β2 .* v .+ (1 - β2) .* (g .* g)
    m̂ = m ./ (1 - β1^t)
    v̂ = v ./ (1 - β2^t)
    θ .-= α .* m̂ ./ (sqrt.(v̂) .+ ϵ)

    (t % 50 == 0) && println("step $t: loss = ", round(occupancy_loss(pieces, pts, θ); digits = 6))
end

step 50: loss = 0.014662
step 100: loss = 0.003488
step 150: loss = 0.002698
step 200: loss = 0.003089
step 250: loss = 0.003689
step 300: loss = 0.002956

assembled = reduce(union, [
    rotate(translate(pieces[i], θ[6i - 5], θ[6i - 4], θ[6i - 3]), rotation_exp(skew(θ[(6i - 2):6i])))
    for i in eachindex(pieces)
])
println("reassembled volume = ", round(volume(surface_mesh(assembled; n=32)); digits=3), "  (want 8)")
vtk(assembled; n=64)

reassembled volume = 7.574 m^3  (want 8)

Zoom Opacity

4. When to use which

assemble (LBFGS + surface) — simplest; fine for near-home pieces. Cheap, but blind to the interior and easily stuck when pieces start scattered.
Recenter → solve → resample — a cheap, deterministic upgrade: the centroid alignment gives LBFGS a reachable start, and the shrinking-box loop refines the boundary. Good when the shuffle is mostly a bulk translation + modest rotations.
Stochastic Adam — the most robust to bad starts (the per-step resampling never lets the gradient go fully stale), at the cost of hand-tuning α and running more steps. Use it when pieces are genuinely scattered or when you want a memory-light, anytime solver.

Both strategies share the same ingredients — the occupancy loss of §1 and the pose/soft-union plumbing — so you can mix them: recenter first, then finish with a few hundred Adam steps on shrinking boxes.

Next steps

Shape Similarity, Registration, and Self-Assembly — the built-in assemble / register and the overlap measures these strategies build on.
Differentiable Simulation and Optimization — the gradient-based driver minimize wraps.