Spaces:
Running on Zero
Running on Zero
| """ | |
| PointNet-based encoder for rigid body representation. | |
| Contains the per-object point cloud encoder (PointNet backbone with spatial | |
| and feature transforms), multi-strategy global pooling, and step-size FiLM | |
| conditioning. | |
| """ | |
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| from typing import Tuple, Optional | |
| class STN3d(nn.Module): | |
| """Spatial Transformer Network for 3D point alignment.""" | |
| def __init__(self, input_dim: int = 3): | |
| super().__init__() | |
| self.input_dim = input_dim | |
| # Convolutional layers | |
| self.conv1 = nn.Conv1d(input_dim, 64, 1, bias=False) | |
| self.conv2 = nn.Conv1d(64, 128, 1, bias=False) | |
| self.conv3 = nn.Conv1d(128, 1024, 1, bias=False) | |
| # Batch normalization | |
| self.bn1 = nn.BatchNorm1d(64) | |
| self.bn2 = nn.BatchNorm1d(128) | |
| self.bn3 = nn.BatchNorm1d(1024) | |
| # Fully connected layers | |
| self.fc1 = nn.Linear(1024, 512, bias=False) | |
| self.fc2 = nn.Linear(512, 256, bias=False) | |
| self.fc3 = nn.Linear(256, input_dim * input_dim, bias=False) | |
| self.bn4 = nn.BatchNorm1d(512) | |
| self.bn5 = nn.BatchNorm1d(256) | |
| def forward(self, x: torch.Tensor) -> torch.Tensor: | |
| """ | |
| Args: | |
| x: Input point cloud (B, input_dim, N) | |
| Returns: | |
| Transformation matrix (B, input_dim, input_dim) | |
| """ | |
| batch_size = x.size(0) | |
| # Extract features (use SiLU for smoother nonlinearity) | |
| x = F.silu(self.bn1(self.conv1(x))) | |
| x = F.silu(self.bn2(self.conv2(x))) | |
| x = F.silu(self.bn3(self.conv3(x))) | |
| # Global feature via max pooling | |
| x = torch.max(x, 2, keepdim=False)[0] | |
| # Transform to matrix | |
| x = F.silu(self.bn4(self.fc1(x))) | |
| x = F.silu(self.bn5(self.fc2(x))) | |
| x = self.fc3(x) | |
| # Add identity matrix for stability | |
| identity = torch.eye(self.input_dim, dtype=x.dtype, device=x.device) | |
| identity = identity.view(1, -1).repeat(batch_size, 1) | |
| x = x + identity | |
| x = x.view(batch_size, self.input_dim, self.input_dim) | |
| return x | |
| def grouped_masked_batch_norm(bn: nn.Module, x: torch.Tensor, mask: torch.Tensor, groups: int) -> torch.Tensor: | |
| """Training-mode BatchNorm over x (G*M, C, N) with separate statistics for each of G groups. | |
| Group g is the contiguous block of rows [g*M, (g+1)*M); its statistics use only the masked-valid | |
| positions (mask (G*M, N), True = valid). This equals G separate calls of the packed masked BatchNorm | |
| (RigidBodyEncoder._bn) on the G blocks, including the running-statistics updates (applied in the | |
| closed form of G sequential momentum updates), without packing: fixed shapes and no host syncs. Under | |
| torch.distributed the per-group sums are all-reduced with an autograd-aware collective, so both the | |
| statistics and their gradients are those of SyncBatchNorm applied per group. Padded positions of the | |
| output are zero, as in the packed path. | |
| Sums run in fp32 with fp64 accumulation (no fp64 copy of x); results match the packed path up to | |
| floating-point rounding (different reduction order).""" | |
| GM, C, N = x.shape | |
| if GM % groups != 0: | |
| raise ValueError(f"grouped_masked_batch_norm: {GM} rows are not divisible into {groups} groups") | |
| M = GM // groups | |
| if mask.shape != (GM, N): | |
| raise ValueError(f"grouped_masked_batch_norm: mask must be {(GM, N)}; got {tuple(mask.shape)}") | |
| m = mask.view(groups, M, 1, N) | |
| xf = x.float().view(groups, M, C, N) | |
| xm = torch.where(m, xf, torch.zeros((), dtype=xf.dtype, device=x.device)) | |
| s1 = xm.sum(dim=(1, 3), dtype=torch.float64) # (G, C) | |
| s2 = (xm * xm).sum(dim=(1, 3), dtype=torch.float64) # (G, C) | |
| cnt = m.sum(dim=(1, 3), dtype=torch.float64) # (G, 1) | |
| if torch.distributed.is_available() and torch.distributed.is_initialized(): | |
| from torch.distributed.nn.functional import all_reduce as _ar_autograd | |
| packed = _ar_autograd(torch.cat([s1, s2, cnt], dim=1)) # (G, 2C+1), gradients all-reduced too | |
| s1, s2, cnt = packed[:, :C], packed[:, C:2 * C], packed[:, 2 * C:] | |
| cnt_c = cnt.clamp_min(1.0) | |
| mean = s1 / cnt_c # (G, C) | |
| var = (s2 / cnt_c - mean * mean).clamp_min(0.0) # biased (normalisation), (G, C) | |
| if bn.track_running_stats and bn.running_mean is not None: | |
| with torch.no_grad(): | |
| mom = bn.momentum if bn.momentum is not None else 0.1 | |
| unbiased = var * (cnt / (cnt - 1.0).clamp_min(1.0)) | |
| # G sequential updates r <- (1-mom) r + mom * stat_g, in closed form | |
| w = mom * (1.0 - mom) ** torch.arange(groups - 1, -1, -1, dtype=torch.float64, device=x.device) | |
| decay = (1.0 - mom) ** groups | |
| rm = bn.running_mean.double() * decay + (w.view(groups, 1) * mean.detach()).sum(0) | |
| rv = bn.running_var.double() * decay + (w.view(groups, 1) * unbiased.detach()).sum(0) | |
| bn.running_mean.copy_(rm.to(bn.running_mean.dtype)) | |
| bn.running_var.copy_(rv.to(bn.running_var.dtype)) | |
| bn.num_batches_tracked.add_(groups) | |
| inv = torch.rsqrt(var.float() + bn.eps).view(groups, 1, C, 1) | |
| y = (xf - mean.float().view(groups, 1, C, 1)) * inv | |
| if bn.affine: | |
| y = y * bn.weight.float().view(1, 1, C, 1) + bn.bias.float().view(1, 1, C, 1) | |
| y = torch.where(m, y, torch.zeros((), dtype=y.dtype, device=y.device)) | |
| return y.view(GM, C, N).to(x.dtype) | |
| def grouped_rows_batch_norm(bn: nn.Module, x: torch.Tensor, group_sizes: list) -> torch.Tensor: | |
| """Training-mode BatchNorm over packed rows x (P, C) with separate statistics for G contiguous row blocks of the | |
| given sizes: the packed counterpart of grouped_masked_batch_norm (same fp64 accumulation, autograd-aware | |
| all-reduce of the per-group sums under torch.distributed, closed-form running-statistics updates of G | |
| sequential momentum steps).""" | |
| G = len(group_sizes) | |
| P, C = x.shape | |
| xf = x.float() | |
| blocks = xf.split(group_sizes, dim=0) | |
| s1 = torch.stack([b.sum(dim=0, dtype=torch.float64) for b in blocks]) # (G, C) | |
| s2 = torch.stack([(b * b).sum(dim=0, dtype=torch.float64) for b in blocks]) # (G, C) | |
| cnt = torch.tensor(group_sizes, dtype=torch.float64, device=x.device).view(G, 1) | |
| if torch.distributed.is_available() and torch.distributed.is_initialized(): | |
| from torch.distributed.nn.functional import all_reduce as _ar_autograd | |
| packed = _ar_autograd(torch.cat([s1, s2, cnt], dim=1)) | |
| s1, s2, cnt = packed[:, :C], packed[:, C:2 * C], packed[:, 2 * C:] | |
| cnt_c = cnt.clamp_min(1.0) | |
| mean = s1 / cnt_c | |
| var = (s2 / cnt_c - mean * mean).clamp_min(0.0) | |
| if bn.track_running_stats and bn.running_mean is not None: | |
| with torch.no_grad(): | |
| mom = bn.momentum if bn.momentum is not None else 0.1 | |
| unbiased = var * (cnt / (cnt - 1.0).clamp_min(1.0)) | |
| w = mom * (1.0 - mom) ** torch.arange(G - 1, -1, -1, dtype=torch.float64, device=x.device) | |
| decay = (1.0 - mom) ** G | |
| rm = bn.running_mean.double() * decay + (w.view(G, 1) * mean.detach()).sum(0) | |
| rv = bn.running_var.double() * decay + (w.view(G, 1) * unbiased.detach()).sum(0) | |
| bn.running_mean.copy_(rm.to(bn.running_mean.dtype)) | |
| bn.running_var.copy_(rv.to(bn.running_var.dtype)) | |
| bn.num_batches_tracked.add_(G) | |
| # Per-block broadcast of the group statistics (the same elementwise ops as indexing them per row, whose | |
| # backward is a heavily duplicated index_put) | |
| inv, mean_f = torch.rsqrt(var.float() + bn.eps), mean.float() # (G, C) | |
| y = torch.cat([(b - mean_f[g]) * inv[g] for g, b in enumerate(blocks)], dim=0) # (P, C) | |
| if bn.affine: | |
| y = y * bn.weight.float() + bn.bias.float() | |
| return y.to(x.dtype) | |
| def hierarchical_subset_size(n_valid: torch.Tensor, level: int, small: str) -> torch.Tensor: | |
| """Vertices pooled by subsampled level `level` (0-based) of an object with n_valid valid vertices: the first | |
| n_valid // 2^(level+1) of its FPS ordering (50/25/12.5% for levels 0/1/2). small = 'all_below_8': a level | |
| with fewer than 8 vertices uses all of them; 'fraction': the stated fraction for every object, at least 1 vertex.""" | |
| k = n_valid // (2 ** (level + 1)) | |
| if small == 'all_below_8': | |
| return torch.where(k < 8, n_valid, k) | |
| if small == 'fraction': | |
| return k.clamp_min(1) | |
| raise ValueError(f"subset_small_objects must be 'all_below_8' or 'fraction', got {small!r}") | |
| class EncoderPack: | |
| """Packed layout of the valid vertices of an encoder call (rows = objects, N = vertex slots): flat indices of the | |
| valid vertices, of every hierarchical level's subset, and the per-group row counts for grouped BatchNorm. | |
| The masks do not change over a batch or a rollout, so the pack is built once (a few host syncs) and reused.""" | |
| def __init__(self, valid_mask: torch.Tensor, vtx_rank: Optional[torch.Tensor], num_levels: int, | |
| subsets: bool, groups: int = 1, small: str = 'all_below_8'): | |
| R, N = valid_mask.shape | |
| if R % groups != 0: | |
| raise ValueError(f"EncoderPack: {R} rows are not divisible into {groups} groups") | |
| self.R, self.N, self.groups = R, N, groups | |
| mask = valid_mask.bool() | |
| self.idx = mask.reshape(-1).nonzero(as_tuple=True)[0] # (P,) sorted | |
| rows_per_group = R // groups | |
| g_of = torch.div(torch.div(self.idx, N, rounding_mode='floor'), rows_per_group, rounding_mode='floor') | |
| self.group_sizes = torch.bincount(g_of, minlength=groups).tolist() if groups > 1 else [int(self.idx.numel())] | |
| self.level_idx, self.level_rows, self.level_masks, self.level_group_sizes = [], [], [], [] | |
| n_valid = mask.sum(dim=1) | |
| # Segment layout for pooling over the packed rows (packed_pooling='segment'): row r's valid vertices are | |
| # the contiguous packed rows of segment r (idx is sorted), lengths = their counts, seg = segment per row. | |
| self.lengths, self.seg = n_valid, torch.div(self.idx, N, rounding_mode='floor') | |
| self.level_lengths, self.level_seg = [], [] | |
| for level in range(num_levels): | |
| if subsets: | |
| if vtx_rank is None: | |
| raise ValueError("EncoderPack: hierarchical subsets need vtx_rank") | |
| k = hierarchical_subset_size(n_valid, level, small) | |
| lmask = mask & (vtx_rank < k.unsqueeze(1)) | |
| else: | |
| lmask = mask | |
| li = lmask.reshape(-1).nonzero(as_tuple=True)[0] # sorted, subset of idx | |
| # positions of the level's vertices inside the packed rows (idx is sorted, li a subset of it) | |
| self.level_rows.append(torch.searchsorted(self.idx, li)) | |
| self.level_idx.append(li) | |
| self.level_masks.append(lmask) | |
| lg = torch.div(torch.div(li, N, rounding_mode='floor'), rows_per_group, rounding_mode='floor') | |
| self.level_group_sizes.append(torch.bincount(lg, minlength=groups).tolist() if groups > 1 | |
| else [int(li.numel())]) | |
| self.level_lengths.append(lmask.sum(dim=1)) | |
| self.level_seg.append(torch.div(li, N, rounding_mode='floor')) | |
| self.mask = mask | |
| class RigidBodyEncoder(nn.Module): | |
| """ | |
| PointNet encoder for rigid body representation. | |
| Follows the original PointNet architecture (STN, feature transform, | |
| shared MLPs), with two pooling variants beyond plain max pooling: | |
| - Enhanced pooling: max + mean + std for richer global features | |
| - Hierarchical pooling: multi-scale aggregation (PointNet++ style) | |
| """ | |
| def __init__( | |
| self, | |
| input_dim: int = 9, # 3 for dist_vec + 3 for velocity + 3 for relative_pos | |
| hidden_dims: list = None, # Will use original PointNet dims by default | |
| output_dim: int = 1024, # Final fused output dim used by downstream modules (MultiBodyEncoder maps to this) | |
| backbone_dim: int = None, # Internal PointNet feature width before pooling/ fusion; defaults to output_dim | |
| use_stn: bool = True, # Enable STN like original | |
| use_batch_norm: bool = True, | |
| feature_transform: bool = True, # Enable feature transform like original | |
| feature_transform_dim: int = 64, # K=64 in original paper | |
| pooling_type: str = 'enhanced', # 'max' | 'enhanced' (max+mean+std) | 'hierarchical' (PointNet++) | |
| num_hierarchical_levels: int = 2, | |
| masked_pooling: bool = False, # Exact set-function pooling / BatchNorm over valid vertices only | |
| freeze_bn_stats: bool = False, # Training uses the stored running statistics (fixed per-channel affine); no re-estimation | |
| height_input: bool = False, # Extra input channel: vertex height above the ground plane (absolute z) | |
| hierarchical_subsets: bool = False, # Levels pool the first 50/25/12.5/... percent of an order-independent FPS ordering | |
| packed_pooling: str = 'scatter', # Packed path only: 'scatter' (pool the padded layout) | 'segment' (pool the packed rows) | |
| subset_small_objects: str = 'all_below_8', # Subset levels below 8 vertices: 'all_below_8' (all vertices) | 'fraction' | |
| ): | |
| super().__init__() | |
| self.input_dim = input_dim | |
| self.output_dim = output_dim | |
| self.backbone_dim = backbone_dim if backbone_dim is not None else output_dim | |
| self.use_stn = use_stn | |
| self.feature_transform = feature_transform | |
| self.feature_transform_dim = feature_transform_dim | |
| self.pooling_type = pooling_type | |
| self.num_hierarchical_levels = num_hierarchical_levels | |
| # Set-function mode: pooling and BatchNorm statistics ignore padded vertices, | |
| # and every hierarchical level pools over the full valid set. | |
| self.masked_pooling = masked_pooling | |
| self.freeze_bn_stats = bool(freeze_bn_stats) | |
| self.height_input = bool(height_input) | |
| self.hierarchical_subsets = bool(hierarchical_subsets) | |
| if packed_pooling not in ('scatter', 'segment'): | |
| raise ValueError(f"packed_pooling must be 'scatter' or 'segment', got {packed_pooling!r}") | |
| self.packed_pooling = packed_pooling | |
| if subset_small_objects not in ('all_below_8', 'fraction'): | |
| raise ValueError(f"subset_small_objects must be 'all_below_8' or 'fraction', got {subset_small_objects!r}") | |
| self.subset_small_objects = subset_small_objects | |
| # Use original PointNet architecture dimensions if not specified: | |
| # conv(64) -> conv(64) -> [feature transform] -> conv(64) -> conv(128) -> conv(1024) | |
| if hidden_dims is None: | |
| hidden_dims = [64, 128, 1024] | |
| # Spatial Transformer Network (for input alignment) | |
| if use_stn: | |
| self.stn = STN3d(input_dim) | |
| # First set of convolutions (before feature transform) | |
| self.conv1 = nn.Conv1d(input_dim, 64, 1, bias=not use_batch_norm) | |
| self.bn1 = nn.BatchNorm1d(64) if use_batch_norm else nn.Identity() | |
| self.conv2 = nn.Conv1d(64, 64, 1, bias=not use_batch_norm) | |
| self.bn2 = nn.BatchNorm1d(64) if use_batch_norm else nn.Identity() | |
| # Feature transform network at 64-dimensional features | |
| if feature_transform: | |
| self.feature_stn = STN3d(feature_transform_dim) | |
| # Second set of convolutions (after feature transform) | |
| self.conv3 = nn.Conv1d(64, 64, 1, bias=not use_batch_norm) | |
| self.bn3 = nn.BatchNorm1d(64) if use_batch_norm else nn.Identity() | |
| self.conv4 = nn.Conv1d(64, 128, 1, bias=not use_batch_norm) | |
| self.bn4 = nn.BatchNorm1d(128) if use_batch_norm else nn.Identity() | |
| self.conv5 = nn.Conv1d(128, self.backbone_dim, 1, bias=not use_batch_norm) | |
| self.bn5 = nn.BatchNorm1d(self.backbone_dim) if use_batch_norm else nn.Identity() | |
| # Enhanced pooling: project (max+mean+std) back to backbone_dim | |
| if pooling_type == 'enhanced': | |
| self.pooling_proj = nn.Sequential( | |
| nn.Linear(self.backbone_dim * 3, self.backbone_dim * 2), | |
| nn.ReLU(), | |
| nn.Linear(self.backbone_dim * 2, self.backbone_dim) | |
| ) | |
| # Hierarchical pooling: multi-scale feature aggregation (PointNet++ inspired) | |
| # Designed to handle variable vertex counts robustly | |
| if pooling_type == 'hierarchical': | |
| self.hierarchical_layers = nn.ModuleList() | |
| for i in range(num_hierarchical_levels): | |
| # Each level: Conv1d for local feature extraction | |
| self.hierarchical_layers.append(nn.Sequential( | |
| nn.Conv1d(self.backbone_dim, self.backbone_dim, 1), | |
| nn.BatchNorm1d(self.backbone_dim) if use_batch_norm else nn.Identity(), | |
| nn.ReLU() | |
| )) | |
| # Project concatenated multi-scale features | |
| # Input: global (backbone_dim*2) + level1 (backbone_dim*2) + ... + levelN (backbone_dim*2) | |
| # Total: backbone_dim * 2 * (1 + num_hierarchical_levels) | |
| proj_input_dim = self.backbone_dim * 2 * (1 + num_hierarchical_levels) | |
| self.hierarchical_proj = nn.Sequential( | |
| nn.Linear(proj_input_dim, self.backbone_dim * 2), | |
| nn.ReLU(), | |
| nn.Dropout(0.1), # Dropout for regularization | |
| nn.Linear(self.backbone_dim * 2, self.backbone_dim) | |
| ) | |
| def _bn(self, bn: nn.Module, x: torch.Tensor, mask: Optional[torch.Tensor], | |
| mask_idx: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, groups: int = 1) -> torch.Tensor: | |
| """BatchNorm over (B, C, N) whose training statistics come from valid vertices only. | |
| The valid vertices are packed into one (1, C, P) tensor and normalised by the stock | |
| (Sync)BatchNorm module, so batch statistics, running statistics and the gradient | |
| through the statistics are exactly those of a BatchNorm applied to the un-padded | |
| set. Padded rows of the output are zero; the masked pooling ignores them. | |
| Evaluation uses the running statistics, a per-channel affine map that is already | |
| independent of padding. | |
| mask_idx: optional mask.nonzero(as_tuple=True), shared by calls that use the same mask. Indexing | |
| with it is what boolean-mask indexing does internally, so results are identical; it only avoids | |
| recomputing the nonzero (a host sync) for the pack and again for the scatter. | |
| groups > 1: the rows are G contiguous blocks (e.g. G time steps batched together) whose statistics | |
| must stay separate; grouped_masked_batch_norm computes them per block without packing.""" | |
| is_bn = isinstance(bn, torch.nn.modules.batchnorm._BatchNorm) | |
| if self.freeze_bn_stats and self.training and is_bn: | |
| # Fine-tuning a converged model: keep the stored statistics (the downstream weights | |
| # were trained against them), so training and evaluation normalise identically and | |
| # padded rows cannot influence anything. The affine parameters remain trainable. | |
| if not getattr(self, '_frozen_bn_logged', False): | |
| self._frozen_bn_logged = True | |
| print(f"[Encoder] BatchNorm statistics FROZEN in training (module={type(bn).__name__}; stored running stats, affine trainable)", flush=True) | |
| return torch.nn.functional.batch_norm(x, bn.running_mean, bn.running_var, bn.weight, bn.bias, False, 0.0, bn.eps) | |
| if not (self.masked_pooling and mask is not None and self.training and is_bn): | |
| return bn(x) | |
| if not getattr(self, '_packed_bn_logged', False): | |
| self._packed_bn_logged = True | |
| print(f"[Encoder] packed masked BatchNorm ACTIVE: module={type(bn).__name__}, valid={int(mask.sum())}/{mask.numel()} vertices", flush=True) | |
| if groups > 1: | |
| return grouped_masked_batch_norm(bn, x, mask, groups) | |
| if mask_idx is None: | |
| mask_idx = mask.nonzero(as_tuple=True) | |
| xt = x.transpose(1, 2) # (B, N, C) | |
| packed = xt[mask_idx] # (P, C), same rows and order as xt[mask] | |
| y = bn(packed.t().unsqueeze(0)).squeeze(0).t() # (P, C) | |
| out = xt.new_zeros(xt.shape) | |
| out[mask_idx] = y.to(out.dtype) | |
| return out.transpose(1, 2) | |
| def _masked_max(f: torch.Tensor, mask: torch.Tensor) -> torch.Tensor: | |
| neg = torch.finfo(f.dtype).min | |
| out = f.masked_fill(~mask.unsqueeze(1), neg).max(dim=2)[0] | |
| return torch.where(mask.any(dim=1, keepdim=True), out, torch.zeros_like(out)) | |
| def _masked_mean(f: torch.Tensor, mask: torch.Tensor) -> torch.Tensor: | |
| # Accumulate in fp32 so the mean does not depend on the padded length under autocast. | |
| m = mask.unsqueeze(1).float() | |
| return ((f.float() * m).sum(dim=2) / m.sum(dim=2).clamp_min(1.0)).to(f.dtype) | |
| def _segment_max_mean(h: torch.Tensor, lengths: torch.Tensor, seg: torch.Tensor): | |
| """_masked_max and _masked_mean of the packed rows h (P, C) without the padded layout: segment r is the | |
| contiguous block of lengths[r] rows (seg = segment of every row). The max has the same value and passes its | |
| gradient to the first maximal row of the segment, which is the index torch.max returns on the padded layout | |
| (the packed order is the slot order); the mean accumulates in fp32. Empty segments give zeros.""" | |
| P = h.shape[0] | |
| if P >= 2 ** 24: | |
| raise ValueError(f"_segment_max_mean: {P} rows exceed the exact fp32 row positions") | |
| with torch.no_grad(): | |
| mx = torch.segment_reduce(h, 'max', lengths=lengths, axis=0, unsafe=True) # (R, C) | |
| pos = torch.arange(P, device=h.device, dtype=torch.float32).unsqueeze(1) | |
| first = torch.segment_reduce(torch.where(h == mx.index_select(0, seg), pos, float(P)), 'min', | |
| lengths=lengths, axis=0, unsafe=True) | |
| first = first.clamp_max(P - 1).long() # (R, C) | |
| mx = torch.where((lengths > 0).unsqueeze(1), h.gather(0, first), torch.zeros((), dtype=h.dtype, device=h.device)) | |
| mean = (torch.segment_reduce(h.float(), 'sum', lengths=lengths, axis=0, unsafe=True) | |
| / lengths.clamp_min(1).unsqueeze(1).float()).to(h.dtype) | |
| return mx, mean | |
| def _bn_rows(self, bn: nn.Module, x: torch.Tensor, group_sizes: list, groups: int) -> torch.Tensor: | |
| """BatchNorm over packed valid rows x (P, C): the statistics of the masked path without packing/scattering.""" | |
| if not isinstance(bn, torch.nn.modules.batchnorm._BatchNorm): | |
| return bn(x) | |
| if self.freeze_bn_stats and self.training: | |
| return torch.nn.functional.batch_norm(x, bn.running_mean, bn.running_var, bn.weight, bn.bias, False, 0.0, bn.eps) | |
| if self.training and groups > 1: | |
| return grouped_rows_batch_norm(bn, x, group_sizes) | |
| return bn(x) | |
| def _rows_linear(conv: nn.Conv1d, x: torch.Tensor) -> torch.Tensor: | |
| return F.linear(x, conv.weight.squeeze(-1), conv.bias) | |
| def _forward_packed(self, x: torch.Tensor, pack: 'EncoderPack', bn_groups: int): | |
| """Encoder over the valid vertices only (model.encoder.packed): every pointwise layer runs on the packed rows, | |
| each hierarchical level only on its own vertex subset, and the features are scattered back (zeros at | |
| padding) for the unchanged masked pooling and for anchor-vertex pooling. Equal to the padded path at every | |
| valid vertex up to floating-point rounding (different GEMM shapes / reduction order).""" | |
| if not (self.masked_pooling and self.pooling_type == 'hierarchical') or self.use_stn or self.feature_transform: | |
| raise ValueError("the packed encoder path needs masked hierarchical pooling without STN / feature transform") | |
| R, N, Cin = x.shape | |
| if (R, N) != (pack.R, pack.N) or bn_groups != pack.groups: | |
| raise ValueError(f"EncoderPack built for {(pack.R, pack.N, pack.groups)}, called with {(R, N, bn_groups)}") | |
| if not getattr(self, '_packed_path_logged', False): | |
| self._packed_path_logged = True | |
| print(f"[Encoder] packed vertex path ACTIVE: {int(pack.idx.numel())}/{R * N} vertex slots valid, " | |
| f"levels {[int(i.numel()) for i in pack.level_idx]}", flush=True) | |
| gs = pack.group_sizes | |
| h = x.reshape(R * N, Cin).index_select(0, pack.idx) | |
| h = F.silu(self._bn_rows(self.bn1, self._rows_linear(self.conv1, h), gs, bn_groups)) | |
| h = F.silu(self._bn_rows(self.bn2, self._rows_linear(self.conv2, h), gs, bn_groups)) | |
| h = F.silu(self._bn_rows(self.bn3, self._rows_linear(self.conv3, h), gs, bn_groups)) | |
| h = F.silu(self._bn_rows(self.bn4, self._rows_linear(self.conv4, h), gs, bn_groups)) | |
| h = self._bn_rows(self.bn5, self._rows_linear(self.conv5, h), gs, bn_groups) # (P, C) | |
| C = h.shape[1] | |
| # Per-vertex features in the padded layout (zeros at padding) for anchor-vertex pooling | |
| per_point_features = h.new_zeros(R * N, C).index_copy(0, pack.idx, h).view(R, N, C).transpose(1, 2) | |
| segment = self.packed_pooling == 'segment' | |
| if segment and not getattr(self, '_segment_pool_logged', False): | |
| self._segment_pool_logged = True | |
| print("[Encoder] packed pooling: segment max/mean over the packed rows (no padded level tensors)", flush=True) | |
| if segment: | |
| feats = [torch.cat(self._segment_max_mean(h, pack.lengths, pack.seg), dim=1)] | |
| else: | |
| mask = pack.mask | |
| feats = [torch.cat([self._masked_max(per_point_features, mask), | |
| self._masked_mean(per_point_features, mask)], dim=1)] | |
| for level, layer in enumerate(self.hierarchical_layers): | |
| hl = self._rows_linear(layer[0], h.index_select(0, pack.level_rows[level])) | |
| hl = layer[2](self._bn_rows(layer[1], hl, pack.level_group_sizes[level], bn_groups)) | |
| if segment: | |
| feats.append(torch.cat(self._segment_max_mean(hl, pack.level_lengths[level], pack.level_seg[level]), | |
| dim=1)) | |
| continue | |
| lf = hl.new_zeros(R * N, C).index_copy(0, pack.level_idx[level], hl).view(R, N, C).transpose(1, 2) | |
| lmask = pack.level_masks[level] | |
| feats.append(torch.cat([self._masked_max(lf, lmask), self._masked_mean(lf, lmask)], dim=1)) | |
| latent = self.hierarchical_proj(torch.cat(feats, dim=1)) | |
| return latent, per_point_features, None | |
| def forward( | |
| self, | |
| dist_vec_norm: torch.Tensor, | |
| velocity_norm: torch.Tensor, | |
| relative_pos: torch.Tensor, | |
| phys_params: Optional[torch.Tensor] = None, | |
| time_step: Optional[torch.Tensor] = None, | |
| valid_mask: Optional[torch.Tensor] = None, | |
| height: Optional[torch.Tensor] = None, | |
| vtx_rank: Optional[torch.Tensor] = None, | |
| bn_groups: int = 1, | |
| pack: Optional['EncoderPack'] = None | |
| ) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]: | |
| """ | |
| Encode rigid body point cloud into latent representation. | |
| Follows the original PointNet architecture. | |
| Args: | |
| dist_vec_norm: Normalized distance vectors (B, N, 3) where N is number of vertices | |
| velocity_norm: Normalized velocities (B, N, 3) | |
| relative_pos: Position relative to reference pose (B, N, 3) | |
| phys_params: Optional physics parameters (B, N, 3) - [mass, friction, restitution] broadcasted to all vertices | |
| time_step: Optional time step size (B, N, 1) - skip_step_size broadcasted to all vertices | |
| valid_mask: Optional vertex validity mask (B, N) - True for valid vertices, False for padding | |
| Returns: | |
| latent: Global latent representation (B, output_dim) | |
| per_point_features: Per-point features before pooling (B, output_dim, N) | |
| trans_feat: Feature transformation matrix if feature_transform=True | |
| """ | |
| B, N, _ = dist_vec_norm.shape | |
| # Concatenate all input features: (B, N, input_dim) | |
| # Base: 9D (dist_vec + velocity + rel_pos) | |
| # +phys_params: +3D = 12D | |
| # +time_step: +2D (s and s^2) = 14D | |
| features = [dist_vec_norm, velocity_norm, relative_pos] | |
| if phys_params is not None: | |
| features.append(phys_params) | |
| if time_step is not None: | |
| features.append(time_step) | |
| if self.height_input: | |
| if height is None: | |
| raise ValueError("RigidBodyEncoder(height_input=True) requires the `height` tensor (B, N, 1)") | |
| features.append(height) # last channel: height above the ground plane | |
| if not getattr(self, '_height_logged', False): | |
| self._height_logged = True | |
| print(f"[Encoder] height-above-ground input channel ACTIVE (input_dim={self.input_dim}, " | |
| f"height range [{height.min().item():.2f}, {height.max().item():.2f}] m)", flush=True) | |
| x = torch.cat(features, dim=-1) | |
| if pack is not None: | |
| return self._forward_packed(x, pack, bn_groups) | |
| # Transpose for Conv1d: (B, N, C) -> (B, C, N) | |
| x = x.transpose(1, 2) | |
| # Apply input spatial transformer if enabled | |
| if self.use_stn: | |
| trans = self.stn(x) | |
| x = x.transpose(2, 1) # (B, C, N) -> (B, N, C) | |
| x = torch.bmm(x, trans) # Apply transformation | |
| x = x.transpose(2, 1) # (B, N, C) -> (B, C, N) | |
| # The five backbone BatchNorms share one vertex mask: compute its index set once (training with | |
| # masked statistics only; evaluation and frozen statistics never pack). | |
| _bn_idx = (valid_mask.nonzero(as_tuple=True) | |
| if (self.masked_pooling and valid_mask is not None and self.training | |
| and not self.freeze_bn_stats and bn_groups == 1) else None) | |
| # First block of convolutions | |
| x = F.silu(self._bn(self.bn1, self.conv1(x), valid_mask, _bn_idx, bn_groups)) | |
| x = F.silu(self._bn(self.bn2, self.conv2(x), valid_mask, _bn_idx, bn_groups)) | |
| # Apply feature transform at 64d features | |
| trans_feat = None | |
| if self.feature_transform: | |
| trans_feat = self.feature_stn(x) | |
| x_transformed = x.transpose(2, 1) # (B, 64, N) -> (B, N, 64) | |
| x_transformed = torch.bmm(x_transformed, trans_feat) # Apply transformation | |
| x = x_transformed.transpose(2, 1) # (B, N, 64) -> (B, 64, N) | |
| # Second block of convolutions (after feature transform) | |
| x = F.silu(self._bn(self.bn3, self.conv3(x), valid_mask, _bn_idx, bn_groups)) | |
| x = F.silu(self._bn(self.bn4, self.conv4(x), valid_mask, _bn_idx, bn_groups)) | |
| x = self._bn(self.bn5, self.conv5(x), valid_mask, _bn_idx, bn_groups) # No ReLU on last layer before pooling | |
| # Store per-point features | |
| per_point_features = x # (B, backbone_dim, N) | |
| # With zero-valued padding, standard pooling is safe: zeros rarely | |
| # dominate the max and only mildly dilute the mean, so valid_mask is | |
| # accepted for interface compatibility but not applied here. | |
| if self.masked_pooling and valid_mask is not None: | |
| _key = '_masked_pool_logged_' + ('train' if self.training else 'eval') | |
| if not getattr(self, _key, False): | |
| setattr(self, _key, True) | |
| print(f"[Encoder] masked set-function pooling ACTIVE ({'train' if self.training else 'eval'} mode): " | |
| f"pooling={self.pooling_type}, levels={getattr(self, 'num_hierarchical_levels', None)}, " | |
| f"valid={int(valid_mask.sum())}/{valid_mask.numel()} vertices, mask dtype={valid_mask.dtype}", flush=True) | |
| mask = valid_mask.bool() | |
| if self.pooling_type == 'max': | |
| latent = self._masked_max(per_point_features, mask) | |
| elif self.pooling_type == 'enhanced': | |
| mean_feat = self._masked_mean(per_point_features, mask) | |
| m = mask.unsqueeze(1).to(per_point_features.dtype) | |
| var = (((per_point_features - mean_feat.unsqueeze(-1)) ** 2) * m).sum(dim=2) / (m.sum(dim=2) - 1).clamp_min(1.0) | |
| combined = torch.cat([self._masked_max(per_point_features, mask), mean_feat, var.clamp_min(1e-12).sqrt()], dim=1) | |
| latent = self.pooling_proj(combined) | |
| elif self.pooling_type == 'hierarchical': | |
| # Every level pools over the full valid set: a set function of the | |
| # vertices, independent of their order, count, and padding. | |
| feats = [torch.cat([self._masked_max(per_point_features, mask), | |
| self._masked_mean(per_point_features, mask)], dim=1)] | |
| if self.hierarchical_subsets: | |
| if vtx_rank is None: | |
| raise ValueError("RigidBodyEncoder(hierarchical_subsets=True) requires `vtx_rank` (B, N)") | |
| n_valid = mask.sum(dim=1) # (B,) | |
| if not getattr(self, '_subsets_logged', False): | |
| self._subsets_logged = True | |
| print(f"[Encoder] hierarchical FPS subsets ACTIVE: levels={len(self.hierarchical_layers)}, " | |
| f"fractions={[1.0 / 2 ** (l + 1) for l in range(len(self.hierarchical_layers))]}, " | |
| f"small objects={self.subset_small_objects}", flush=True) | |
| for level, layer in enumerate(self.hierarchical_layers): | |
| if self.hierarchical_subsets: | |
| # n // 2^(level+1) vertices; small objects per subset_small_objects | |
| k = hierarchical_subset_size(n_valid, level, self.subset_small_objects) | |
| level_mask = mask & (vtx_rank < k.unsqueeze(1)) | |
| else: | |
| level_mask = mask | |
| h = layer[0](per_point_features) | |
| h = layer[2](self._bn(layer[1], h, level_mask, None, bn_groups)) | |
| feats.append(torch.cat([self._masked_max(h, level_mask), self._masked_mean(h, level_mask)], dim=1)) | |
| latent = self.hierarchical_proj(torch.cat(feats, dim=1)) | |
| else: | |
| raise ValueError(f"Unknown pooling_type: {self.pooling_type}") | |
| return latent, per_point_features, trans_feat | |
| # Unmasked pooling over all slots (PADDING_VALUE=0). | |
| if self.pooling_type == 'max': | |
| latent = torch.max(per_point_features, dim=2)[0] # (B, backbone_dim) | |
| elif self.pooling_type == 'enhanced': | |
| max_feat = torch.max(per_point_features, dim=2)[0] # (B, backbone_dim) | |
| mean_feat = torch.mean(per_point_features, dim=2) # (B, backbone_dim) | |
| std_feat = torch.std(per_point_features, dim=2) # (B, backbone_dim) | |
| combined = torch.cat([max_feat, mean_feat, std_feat], dim=1) # (B, backbone_dim*3) | |
| latent = self.pooling_proj(combined) # (B, backbone_dim) | |
| elif self.pooling_type == 'hierarchical': | |
| # Hierarchical: multi-scale feature aggregation (PointNet++ inspired) | |
| global_max = torch.max(per_point_features, dim=2)[0] # (B, backbone_dim) | |
| global_mean = torch.mean(per_point_features, dim=2) # (B, backbone_dim) | |
| global_feat = torch.cat([global_max, global_mean], dim=1) # (B, backbone_dim*2) | |
| local_features = [global_feat] | |
| for level, layer in enumerate(self.hierarchical_layers): | |
| stride = 2 ** (level + 1) | |
| sampled_feat = per_point_features[:, :, ::stride] # (B, backbone_dim, N//stride) | |
| if sampled_feat.shape[2] < 8: | |
| sampled_feat = per_point_features | |
| local_feat = layer(sampled_feat) # (B, backbone_dim, num_sampled) | |
| local_max = torch.max(local_feat, dim=2)[0] # (B, backbone_dim) | |
| local_mean = torch.mean(local_feat, dim=2) # (B, backbone_dim) | |
| local_combined = torch.cat([local_max, local_mean], dim=1) # (B, backbone_dim*2) | |
| local_features.append(local_combined) | |
| all_features = torch.cat(local_features, dim=1) | |
| latent = self.hierarchical_proj(all_features) # (B, backbone_dim) | |
| else: | |
| raise ValueError(f"Unknown pooling_type: {self.pooling_type}") | |
| return latent, per_point_features, trans_feat | |
| class StepSizeFiLM(nn.Module): | |
| """ | |
| Feature-wise Linear Modulation (FiLM) conditioned on step-size. | |
| Generates per-feature scale (γ) and shift (β) from [s, s²]: | |
| output = features * (1 + γ) + β | |
| This gives the step size a stronger influence than concatenating | |
| [s, s²] as two extra dimensions of the encoder input. | |
| Reference: Perez et al., "FiLM: Visual Reasoning with a General | |
| Conditioning Layer", AAAI 2018. | |
| """ | |
| def __init__(self, time_dim: int = 2, feature_dim: int = 768): | |
| super().__init__() | |
| self.film_net = nn.Sequential( | |
| nn.Linear(time_dim, 256), | |
| nn.SiLU(), | |
| nn.Linear(256, feature_dim * 2) # γ and β | |
| ) | |
| # Zero-init the head so γ = β = 0 and FiLM starts as the identity | |
| nn.init.zeros_(self.film_net[-1].weight) | |
| nn.init.zeros_(self.film_net[-1].bias) | |
| def forward(self, features: torch.Tensor, time_features: torch.Tensor) -> torch.Tensor: | |
| """ | |
| Args: | |
| features: (B, Nobj, D) or (B*Nobj, Nq, D) — encoder/predictor features | |
| time_features: (B, 2) — [s, s²] per sample | |
| Returns: | |
| Modulated features, same shape as input | |
| """ | |
| gamma_beta = self.film_net(time_features) # (B, 2*D) | |
| D = features.shape[-1] | |
| gamma = gamma_beta[..., :D] # (B, D) | |
| beta = gamma_beta[..., D:] # (B, D) | |
| # Broadcast over spatial dimensions (Nobj or Nq) | |
| ndim_diff = features.dim() - gamma.dim() | |
| for _ in range(ndim_diff): | |
| gamma = gamma.unsqueeze(-2) # (B, 1, D) or (B, 1, 1, D) | |
| beta = beta.unsqueeze(-2) | |
| return features * (1.0 + gamma) + beta | |