Spaces:
Running on Zero
Running on Zero
File size: 7,939 Bytes
632d805 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 | """
Differentiable intersection handling.
Resolves object-ground and object-object penetrations by translating whole
objects, with a straight-through estimator (STE) keeping the correction
differentiable. The correction formulation follows UNIC (arXiv:2603.25580,
Sec. 3.3), Eq. 8:
v_corrected = v + detach(target_position - v)
Forward: moves v to target_position
Backward: gradients pass through to v unchanged
Whereas UNIC drags individual garment vertices, rigid objects are translated
as a whole:
- Ground plane: shift the object upward when its lowest vertex falls
below ground_z + buffer
- Object-object: push objects apart when their bounding spheres overlap
"""
import torch
import torch.nn.functional as F
import logging
logger = logging.getLogger(__name__)
def resolve_ground_penetration(
positions: torch.Tensor,
obj_mask: torch.Tensor,
vtx_mask: torch.Tensor,
ground_z: float = 0.0,
buffer: float = 0.005,
) -> torch.Tensor:
"""
Resolve object-ground penetrations by translating objects upward.
For each valid object, find the LOWEST valid vertex. If min_z < ground_z + buffer,
shift ALL vertices of that object upward by the deficit.
This guarantees no valid vertex penetrates the ground, while preserving rigidity.
Uses STE trick: correction is detached so gradients flow through original positions.
Args:
positions: (B, Nobj, Nv, 3) predicted vertex positions
obj_mask: (B, Nobj) boolean mask for valid objects
vtx_mask: (B, Nobj, Nv) boolean mask for valid vertices
ground_z: ground plane height (z-up coordinate system)
buffer: relaxation buffer above ground plane
Returns:
corrected_positions: (B, Nobj, Nv, 3) with ground penetrations resolved
"""
B, Nobj, Nv, _ = positions.shape
target_z = ground_z + buffer
# Find per-object minimum z among valid vertices
# Set invalid vertices to +inf so they don't affect min
z_vals = positions[..., 2] # (B, Nobj, Nv)
z_vals_masked = z_vals.masked_fill(~vtx_mask, float('inf')) # invalid vertices -> +inf
min_z, _ = z_vals_masked.min(dim=2) # (B, Nobj) lowest valid vertex z per object
# Find objects that penetrate the ground
penetration_depth = target_z - min_z # positive means below ground
# Padding objects have min_z=+inf, giving depth=-inf; clamp to >= 0 before
# the expand below so 0*(-inf) can never produce NaN
penetration_depth = penetration_depth.clamp(min=0.0)
is_penetrating = (penetration_depth > 0) & obj_mask # (B, Nobj)
# No early exit on "nothing penetrates": the correction below is identically
# zero in that case, and the shortcut would cost a host sync per call.
# Compute correction: shift entire object upward by penetration_depth
correction = torch.zeros_like(positions) # (B, Nobj, Nv, 3)
# Only z-component needs correction
correction[..., 2] = penetration_depth.unsqueeze(-1).expand(-1, -1, Nv) # (B, Nobj, Nv)
# Zero out correction for non-penetrating objects
correction = correction * is_penetrating.float().unsqueeze(-1).unsqueeze(-1) # (B, Nobj, Nv, 3)
# STE trick (UNIC Eq. 8): forward moves position, backward passes gradient through
corrected = positions + correction.detach()
return corrected
def resolve_object_penetration(
positions: torch.Tensor,
obj_mask: torch.Tensor,
vtx_mask: torch.Tensor,
buffer: float = 0.005,
) -> torch.Tensor:
"""
Resolve object-object penetrations using bounding sphere overlap detection.
For each pair of valid objects, compute COM and bounding sphere radius.
If spheres overlap, push both objects apart along the COM-COM direction.
Uses STE trick for differentiability.
Args:
positions: (B, Nobj, Nv, 3) predicted vertex positions
obj_mask: (B, Nobj) boolean mask for valid objects
vtx_mask: (B, Nobj, Nv) boolean mask for valid vertices
buffer: relaxation buffer between objects
Returns:
corrected_positions: (B, Nobj, Nv, 3) with object penetrations resolved
"""
B, Nobj, Nv, _ = positions.shape
if Nobj < 2:
return positions
# Compute per-object COM and bounding sphere radius
vtx_weight = vtx_mask.float().unsqueeze(-1) # (B, Nobj, Nv, 1)
num_valid = vtx_weight.sum(dim=2).clamp(min=1) # (B, Nobj, 1)
com = (positions * vtx_weight).sum(dim=2) / num_valid # (B, Nobj, 3)
# Bounding sphere radius: max distance from COM to any valid vertex
diff_to_com = positions - com.unsqueeze(2) # (B, Nobj, Nv, 3)
dist_to_com = diff_to_com.norm(dim=-1) # (B, Nobj, Nv)
# Mask invalid vertices with 0 distance
dist_to_com = dist_to_com * vtx_mask.float() # (B, Nobj, Nv)
radius = dist_to_com.max(dim=2).values # (B, Nobj)
# Accumulate corrections per object
correction = torch.zeros(B, Nobj, 3, device=positions.device, dtype=positions.dtype)
# Check all pairs (i, j) where i < j
for i in range(Nobj):
for j in range(i + 1, Nobj):
# Both objects must be valid
pair_valid = obj_mask[:, i] & obj_mask[:, j] # (B,)
if not pair_valid.any():
continue
# COM distance
com_diff = com[:, j] - com[:, i] # (B, 3)
com_dist = com_diff.norm(dim=-1).clamp(min=1e-8) # (B,)
# Minimum allowed distance
min_dist = radius[:, i] + radius[:, j] + buffer # (B,)
# Overlap amount
overlap = min_dist - com_dist # (B,) positive means penetrating
is_overlapping = (overlap > 0) & pair_valid # (B,)
if not is_overlapping.any():
continue
# Push direction: from i to j (normalized)
push_dir = com_diff / com_dist.unsqueeze(-1) # (B, 3)
# Each object gets pushed by half the overlap
push_amount = (overlap / 2).clamp(min=0) * is_overlapping.float() # (B,)
correction[:, i] -= push_amount.unsqueeze(-1) * push_dir # push i away from j
correction[:, j] += push_amount.unsqueeze(-1) * push_dir # push j away from i
if (correction.abs() < 1e-10).all():
return positions
# Expand correction to per-vertex: (B, Nobj, 3) -> (B, Nobj, Nv, 3)
correction_expanded = correction.unsqueeze(2).expand(-1, -1, Nv, -1)
# STE trick
corrected = positions + correction_expanded.detach()
return corrected
def handle_intersections(
positions: torch.Tensor,
obj_mask: torch.Tensor,
vtx_mask: torch.Tensor,
config: dict,
) -> torch.Tensor:
"""
Main entry point: apply all intersection resolution steps.
Args:
positions: (B, Nobj, Nv, 3) predicted vertex positions
obj_mask: (B, Nobj) boolean mask
vtx_mask: (B, Nobj, Nv) boolean mask for valid vertices
config: intersection_handling config dict
Returns:
corrected positions: (B, Nobj, Nv, 3)
"""
if not config.get('enable', False):
return positions
ground_cfg = config.get('ground', {})
object_cfg = config.get('object', {})
# Step 1: Resolve ground penetrations
if ground_cfg.get('enable', True):
positions = resolve_ground_penetration(
positions, obj_mask, vtx_mask,
ground_z=ground_cfg.get('z', 0.0),
buffer=ground_cfg.get('buffer', 0.005),
)
# Step 2: Resolve object-object penetrations
if object_cfg.get('enable', False):
positions = resolve_object_penetration(
positions, obj_mask, vtx_mask,
buffer=object_cfg.get('buffer', 0.005),
)
return positions
|