"""Batched boundary prop sampling for rendering-only instancing.
``PropSampler`` resamples one track boundary (inner or outer) at a user-set
arc-length spacing and emits per-prop instancing poses into a preallocated
:class:`PropSet`. Two modes:
- ``"points"``: one pose per sample on the curve (cones, poles, markers).
- ``"segments"``: one pose per chord between consecutive samples (wall
pieces): position = chord midpoint, yaw = chord direction (plan-view), length
= chord length, so a unit-length wall scaled by ``length`` tiles the ring.
On a lifted (non-flat ``z_profile``) track, ``length`` is the 3D chord
length between the two (possibly different-altitude) samples, not its planar
projection — see ``PropSet.length``.
Spacing is snapped per env — ``n = clamp(round(perimeter/spacing), 3,
max_props)`` at effective step ``perimeter/n`` — so the closed ring has no
seam gap or doubled prop. Props are NOT colliders (see ``track_gen.collision``
for out-of-bounds queries); this is the rendering/instancing complement.
Layout follows the package conventions: flat ``[E * max_props]`` wp.arrays,
NaN-padded past ``count[e]``, in-place reuse of the output ``PropSet`` across
``sample()`` calls (use ``clone()`` for snapshots), and no host syncs while
capturing is enabled (``track_gen.set_capturing``) so ``sample()`` is
CUDA-graph capturable.
"""
from __future__ import annotations
from dataclasses import dataclass
import numpy as np
import warp as wp
from .collision_geom import _safe_normalize2
from .runtime import _init, _sync
from .types import Track
[docs]
@dataclass
class PropSet:
"""Batched prop poses, flat ``[E * max_props]`` per pose field.
.. warning::
``PropSampler.sample()`` returns the SAME ``PropSet`` instance on
every call and overwrites its buffers in place. Call ``clone()`` for
a fully-owned snapshot.
Attributes
----------
position : wp.array
``vec3f`` pose positions on the (possibly lifted) boundary: on-curve
sample (points mode) or chord midpoint (segments mode); z tracks the
boundary elevation (0 on flat tracks). NaN past ``count[e]``.
tangent : wp.array
``vec2f`` unit directions: curve tangent at the sample (points) or
chord direction (segments). NaN past ``count[e]``.
yaw : wp.array
``float32`` heading, ``atan2(tangent.y, tangent.x)``.
length : wp.array
``float32`` per-prop extent: effective arc step (points mode) or
chord length (segments mode). Both are measured along the (possibly
lifted) boundary polyline: on a flat track this is a planar distance,
and on a non-flat ``z_profile`` track it is the true 3D distance (the
3D chord, in segments mode) between samples that may sit at different
altitudes.
count : wp.array
``[E]`` ``int32`` real prop counts (0 for degenerate envs). Meaningful
only for envs with ``valid[e] == 1``.
truncated : wp.array
``[E]`` ``int32`` — 1 if ``max_props`` clipped this env's ring (the
ring still closes, at a coarser effective spacing).
step : wp.array
``[E]`` ``float32`` effective arc spacing ``perimeter / count``.
"""
position: wp.array
tangent: wp.array
yaw: wp.array
length: wp.array
count: wp.array
truncated: wp.array
step: wp.array
[docs]
def clone(self) -> "PropSet":
"""Return a deep copy whose Warp buffers do not alias this set."""
return PropSet(
position=wp.clone(self.position),
tangent=wp.clone(self.tangent),
yaw=wp.clone(self.yaw),
length=wp.clone(self.length),
count=wp.clone(self.count),
truncated=wp.clone(self.truncated),
step=wp.clone(self.step),
)
@wp.kernel
def _scan_boundary_k(
points: wp.array(dtype=wp.vec3f),
count: wp.array(dtype=wp.int32),
n_max: int,
spacing: float,
max_props: int,
cum: wp.array(dtype=wp.float32),
out_count: wp.array(dtype=wp.int32),
out_step: wp.array(dtype=wp.float32),
out_truncated: wp.array(dtype=wp.int32),
):
"""Also launched by ``_src/checkpoints.py`` (CheckpointSampler) — keep the
signature and snap semantics in sync. Arc lengths are the true 3D segment
lengths of the (possibly lifted) polyline; on the flat z = 0 pipeline this
is identical to the xy length (z contributes 0 to every sum)."""
e = wp.tid()
m = count[e]
if m > n_max:
m = n_max
if m < 3:
out_count[e] = 0
out_step[e] = wp.nan
out_truncated[e] = 0
return
base = e * n_max
s = float(0.0)
cum[base] = 0.0
prev = points[base]
for i in range(1, m):
p = points[base + i]
d = p - prev
s = s + wp.length(d)
cum[base + i] = s
prev = p
dc = points[base] - prev # closing edge back to point 0
perim = s + wp.length(dc)
# Snap in float first: float->int32 conversion of out-of-range values
# (absurdly small spacing) differs between CPU (INT_MIN) and CUDA
# (saturating), so clamp against max_props before the cast.
nf = wp.round(perim / spacing)
trunc = int(0)
n = int(0)
if nf > float(max_props):
n = max_props
trunc = int(1)
else:
n = int(nf)
if n < 3:
n = 3
out_count[e] = n
out_step[e] = perim / float(n)
out_truncated[e] = trunc
@wp.func
def _seg_param_at_arc(cum: wp.array(dtype=wp.float32), base: int, m: int,
perim: float, s: float) -> wp.vec2f:
# (segment index as float, lerp t) at arc position s — same binary
# search contract as checkpoints._seg_at_arc.
lo = int(0)
hi = m - 1
while lo < hi:
mid = (lo + hi + 1) // 2
if cum[base + mid] <= s:
lo = mid
else:
hi = mid - 1
i = lo
seg_end = perim
if i + 1 < m:
seg_end = cum[base + i + 1]
seg_start = cum[base + i]
denom = seg_end - seg_start
t = 0.0
if denom > 1.0e-12:
t = wp.clamp((s - seg_start) / denom, 0.0, 1.0)
return wp.vec2f(float(i), t)
@wp.kernel
def _place_props_k(
points: wp.array(dtype=wp.vec3f),
count: wp.array(dtype=wp.int32),
n_max: int,
cum: wp.array(dtype=wp.float32),
prop_count: wp.array(dtype=wp.int32),
prop_step: wp.array(dtype=wp.float32),
max_props: int,
mode: int, # 0 = points, 1 = segments
out_position: wp.array(dtype=wp.vec3f),
out_tangent: wp.array(dtype=wp.vec2f),
out_yaw: wp.array(dtype=wp.float32),
out_length: wp.array(dtype=wp.float32),
):
t = wp.tid()
e = t // max_props
k = t - e * max_props
n = prop_count[e]
if k >= n:
out_position[t] = wp.vec3f(wp.nan, wp.nan, wp.nan)
out_tangent[t] = wp.vec2f(wp.nan, wp.nan)
out_yaw[t] = wp.nan
out_length[t] = wp.nan
return
m = count[e]
if m > n_max:
m = n_max
base = e * n_max
step = prop_step[e]
perim = step * float(n)
if mode == 0:
it = _seg_param_at_arc(cum, base, m, perim, float(k) * step)
i = int(it[0])
tt = it[1]
j = i + 1
if j == m:
j = 0
a = points[base + i]
b = points[base + j]
p = a + (b - a) * tt # lifted 3D sample position
d = b - a
tang = _safe_normalize2(wp.vec2f(d[0], d[1])) # planar pose direction
out_position[t] = p
out_tangent[t] = tang
out_yaw[t] = wp.atan2(tang[1], tang[0])
out_length[t] = step
else:
s0 = float(k) * step
s1 = 0.0 # slot n-1 chords back to sample 0: the ring closes
if k + 1 < n:
s1 = float(k + 1) * step
it0 = _seg_param_at_arc(cum, base, m, perim, s0)
it1 = _seg_param_at_arc(cum, base, m, perim, s1)
i0 = int(it0[0])
j0 = i0 + 1
if j0 == m:
j0 = 0
i1 = int(it1[0])
j1 = i1 + 1
if j1 == m:
j1 = 0
p0 = points[base + i0] + (points[base + j0] - points[base + i0]) * it0[1]
p1 = points[base + i1] + (points[base + j1] - points[base + i1]) * it1[1]
chord = p1 - p0 # 3D chord
tang = _safe_normalize2(wp.vec2f(chord[0], chord[1])) # planar yaw
out_position[t] = (p0 + p1) * 0.5
out_tangent[t] = tang
out_yaw[t] = wp.atan2(tang[1], tang[0])
out_length[t] = wp.length(chord) # 3D chord length (3D-spacing consistent)
[docs]
class PropSampler:
"""Resample one track boundary into instancing poses at a set spacing.
One sampler binds one :class:`Track`, one boundary curve, one mode, and
one spacing, so all output shapes are fixed and ``sample()`` is
CUDA-graph capturable (allocation-free; host-sync-free while capturing
is enabled via ``track_gen.set_capturing``). Because
``TrackGenerator.generate()`` overwrites its ``Track`` buffers in place,
every ``sample()`` reads the CURRENT batch — no rebind after
regeneration. Props are rendering-only; they are not colliders. Results
for envs with ``valid[e] == 0`` are undefined — an invalid track can
still have >= 3 boundary points and will yield well-formed-looking props
with no NaN signal; callers must gate on ``Track.valid``, as everywhere
in the library.
``PropSet.truncated``/``PropSet.step`` live on the RESULT because
``PropSampler`` is the ONLY producer of a ``PropSet`` — unlike
``CheckpointSet``, which is also produced zero-copy by
``CheckpointSet.from_gates`` (no sampling, so no truncation/step to
report); ``track_gen.checkpoints.CheckpointSampler`` keeps its analogous
diagnostics on the sampler instead, to avoid meaningless fields on
gate-derived checkpoint sets.
"""
[docs]
def __init__(self, track: Track, spacing: float, boundary: str = "outer",
mode: str = "points", max_props: "int | None" = None) -> None:
"""Bind to a :class:`Track` boundary and allocate the result buffers.
Args:
track: the bound track batch; ``sample()`` reads its buffers
directly on every call (no rebind needed after
``generate()``).
spacing: target arc-length spacing between props (> 0); snapped
per env so each ring closes without a seam (see module
docstring).
boundary: ``"inner"`` or ``"outer"`` (default) — which track
boundary curve to resample.
mode: ``"points"`` (default, one pose per sample) or
``"segments"`` (one pose per chord between samples).
max_props: buffer capacity per env; ``None`` (default) derives
it from the CURRENT batch (see :meth:`_derive_max_props`).
Must be >= 3 when given explicitly.
"""
_init()
if not (float(spacing) > 0.0):
raise ValueError(f"spacing must be > 0, got {spacing!r}")
if boundary not in ("inner", "outer"):
raise ValueError(
f"boundary must be one of {{'inner', 'outer'}}, got {boundary!r}")
if mode not in ("points", "segments"):
raise ValueError(
f"mode must be one of {{'points', 'segments'}}, got {mode!r}")
if max_props is not None and int(max_props) < 3:
raise ValueError(
f"max_props must be >= 3 (or None for auto), got {max_props!r}")
E = int(track.count.shape[0])
stride = int(track.outer.shape[0])
if E < 1 or stride % E != 0:
raise ValueError(
f"track batch layout invalid: outer has {stride} slots for {E} envs")
self._track = track
self._spacing = float(spacing)
self._boundary = boundary
self._mode = mode
self._mode_int = 0 if mode == "points" else 1
self._E = E
self._n_max = stride // E
self._points = track.inner if boundary == "inner" else track.outer
self._device = str(track.outer.device)
if max_props is None:
max_props = self._derive_max_props()
self._M = int(max_props)
dev = self._device
n = E * self._M
self._cum = wp.zeros(E * self._n_max, dtype=wp.float32, device=dev)
self._props = PropSet(
position=wp.zeros(n, dtype=wp.vec3f, device=dev),
tangent=wp.zeros(n, dtype=wp.vec2f, device=dev),
yaw=wp.zeros(n, dtype=wp.float32, device=dev),
length=wp.zeros(n, dtype=wp.float32, device=dev),
count=wp.zeros(E, dtype=wp.int32, device=dev),
truncated=wp.zeros(E, dtype=wp.int32, device=dev),
step=wp.zeros(E, dtype=wp.float32, device=dev),
)
def _derive_max_props(self) -> int:
"""Slots to cover the batch bound NOW: ceil(1.5 * max perimeter / spacing).
Host-side readback, construction time only. The 1.5x headroom absorbs
longer boundaries from future regenerations; rings that still exceed
it are truncated and flagged (``PropSet.truncated``).
"""
E, n_max = self._E, self._n_max
pts = self._points.numpy().reshape(E, n_max, 3)
counts = self._track.count.numpy()
valid = self._track.valid.numpy()
best = 0.0
for e in range(E):
m = int(min(counts[e], n_max))
if valid[e] == 0 or m < 3:
continue
poly = pts[e, :m]
seg = np.linalg.norm(np.roll(poly, -1, axis=0) - poly, axis=1)
best = max(best, float(seg.sum()))
if best <= 0.0:
raise ValueError(
"max_props=None needs at least one valid env in the bound track "
"batch to derive a buffer size; pass max_props explicitly")
return max(3, int(np.ceil(1.5 * best / self._spacing)))
[docs]
def sample(self) -> PropSet:
"""Resample the bound boundary; returns the preallocated :class:`PropSet`.
Two kernel launches (scan + place); reads the Track buffers as they
are NOW, so it reflects the latest ``generate()`` with no rebind.
"""
t = self._track
p = self._props
wp.launch(
_scan_boundary_k, dim=self._E,
inputs=[self._points, t.count, self._n_max, self._spacing, self._M,
self._cum, p.count, p.step, p.truncated],
device=self._device,
)
wp.launch(
_place_props_k, dim=self._E * self._M,
inputs=[self._points, t.count, self._n_max, self._cum,
p.count, p.step, self._M, self._mode_int,
p.position, p.tangent, p.yaw, p.length],
device=self._device,
)
_sync(self._device)
return p