""" Prediction heads for the dynamics model: plain MLP, multi-scale, and residual variants. """ import torch import torch.nn as nn import torch.nn.functional as F from typing import List, Optional class DynamicsPredictor(nn.Module): """ MLP predictor that takes state features and query points to predict dynamics. """ def __init__( self, state_dim: int = 512, query_dim: int = 3, # 3D position hidden_dims: List[int] = [256, 128, 64], output_dim: int = 3, # 3D acceleration/displacement activation: str = 'relu', use_batch_norm: bool = False, dropout: float = 0.1 ): super().__init__() self.state_dim = state_dim self.query_dim = query_dim self.output_dim = output_dim # Build MLP layers layers = [] input_dim = state_dim + query_dim for hidden_dim in hidden_dims: layers.append(nn.Linear(input_dim, hidden_dim)) if use_batch_norm: layers.append(nn.BatchNorm1d(hidden_dim)) if activation == 'relu': layers.append(nn.SiLU()) elif activation == 'gelu': layers.append(nn.SiLU()) elif activation == 'tanh': layers.append(nn.Tanh()) else: raise ValueError(f"Unknown activation: {activation}") if dropout > 0: layers.append(nn.Dropout(dropout)) input_dim = hidden_dim # Output layer layers.append(nn.Linear(hidden_dims[-1], output_dim)) layers.append(nn.Tanh()) # bounded output self.output_scale = nn.Parameter(torch.ones(1) * 1.0) self.mlp = nn.Sequential(*layers) def forward( self, state_features: torch.Tensor, query_points: torch.Tensor ) -> torch.Tensor: """ Predict dynamics for query points given state features. Args: state_features: Enhanced state features (B, N_state, D_state) or (B, D_state) query_points: Query vertex positions (B, N_query, 3) Returns: Predicted accelerations/displacements (B, N_query, 3) """ B = query_points.shape[0] N_query = query_points.shape[1] # If state_features has sequence dimension, pool it if state_features.dim() == 3: # Global max pooling over state tokens state_features = torch.max(state_features, dim=1)[0] # (B, D_state) # Expand state features to match number of query points state_features_expanded = state_features.unsqueeze(1).expand(B, N_query, -1) # Concatenate state features with query points combined = torch.cat([state_features_expanded, query_points], dim=-1) # Reshape for batch processing through MLP combined_flat = combined.view(B * N_query, -1) # Apply MLP output_flat = self.mlp(combined_flat) * self.output_scale # Reshape back output = output_flat.view(B, N_query, self.output_dim) return output class MultiScalePredictor(nn.Module): """ Multi-scale predictor that processes features at different resolutions. """ def __init__( self, state_dim: int = 512, query_dim: int = 3, scales: List[int] = [64, 128, 256], output_dim: int = 3, activation: str = 'relu', dropout: float = 0.1 ): super().__init__() self.scales = scales # Create predictors for each scale self.scale_predictors = nn.ModuleList([ DynamicsPredictor( state_dim=state_dim, query_dim=query_dim, hidden_dims=[scale, scale // 2], output_dim=output_dim, activation=activation, dropout=dropout ) for scale in scales ]) # Fusion layer self.fusion = nn.Linear(output_dim * len(scales), output_dim) def forward( self, state_features: torch.Tensor, query_points: torch.Tensor ) -> torch.Tensor: """ Multi-scale prediction with fusion. """ predictions = [] for predictor in self.scale_predictors: pred = predictor(state_features, query_points) predictions.append(pred) # Concatenate all scale predictions combined = torch.cat(predictions, dim=-1) # Fuse predictions output = self.fusion(combined) return output class ResidualPredictor(nn.Module): """ Predictor with residual connections for better gradient flow. """ def __init__( self, state_dim: int = 512, query_dim: int = 3, hidden_dim: int = 256, num_layers: int = 3, output_dim: int = 3, activation: str = 'relu', dropout: float = 0.1 ): super().__init__() self.input_projection = nn.Linear(state_dim + query_dim, hidden_dim) # Residual blocks self.res_blocks = nn.ModuleList([ ResidualBlock(hidden_dim, activation, dropout) for _ in range(num_layers) ]) self.output_projection = nn.Linear(hidden_dim, output_dim) def forward( self, state_features: torch.Tensor, query_points: torch.Tensor ) -> torch.Tensor: """Forward pass with residual connections.""" B = query_points.shape[0] N_query = query_points.shape[1] # Pool state features if needed if state_features.dim() == 3: state_features = torch.max(state_features, dim=1)[0] # Expand and concatenate state_expanded = state_features.unsqueeze(1).expand(B, N_query, -1) combined = torch.cat([state_expanded, query_points], dim=-1) # Initial projection x = self.input_projection(combined) # Apply residual blocks for block in self.res_blocks: x = block(x) # Output projection output = self.output_projection(x) return output class ResidualBlock(nn.Module): """Single residual block.""" def __init__(self, hidden_dim: int, activation: str = 'relu', dropout: float = 0.1): super().__init__() self.fc1 = nn.Linear(hidden_dim, hidden_dim) self.fc2 = nn.Linear(hidden_dim, hidden_dim) if activation == 'relu': self.activation = nn.ReLU() elif activation == 'gelu': self.activation = nn.GELU() else: self.activation = nn.ReLU() self.dropout = nn.Dropout(dropout) self.norm1 = nn.LayerNorm(hidden_dim) self.norm2 = nn.LayerNorm(hidden_dim) def forward(self, x): residual = x x = self.norm1(x) x = self.fc1(x) x = self.activation(x) x = self.dropout(x) x = self.norm2(x) x = self.fc2(x) x = self.dropout(x) return x + residual if __name__ == "__main__": """Test predictor modules.""" print("Testing DynamicsPredictor...") predictor = DynamicsPredictor( state_dim=512, query_dim=3, hidden_dims=[256, 128, 64], output_dim=3 ) # Test inputs batch_size = 4 state_features = torch.randn(batch_size, 10, 512) # With sequence dimension query_points = torch.randn(batch_size, 8, 3) # 8 vertices output = predictor(state_features, query_points) print(f"Output shape: {output.shape}") # Should be (4, 8, 3) print("\nTesting MultiScalePredictor...") multi_predictor = MultiScalePredictor(state_dim=512, query_dim=3) output_multi = multi_predictor(state_features, query_points) print(f"Multi-scale output shape: {output_multi.shape}") # (4, 8, 3) print("\nTesting ResidualPredictor...") res_predictor = ResidualPredictor(state_dim=512, query_dim=3) output_res = res_predictor(state_features, query_points) print(f"Residual output shape: {output_res.shape}") # (4, 8, 3) print("\nPredictor tests successful!")