Spaces:
Sleeping
Sleeping
| #!/usr/bin/env python | |
| # -*- coding: utf-8 -*- | |
| # File: physics_model.py | |
| # | |
| # Physics-aware network components for the HJ-Communicative RL study. | |
| # | |
| # Contents | |
| # 1. PhysicsEncoder -- small CNN that ingests the extra physics channels. | |
| # 2. MultiHeadCommNet -- the PDF's "Step 8": replaces CommNet's mean/single-head | |
| # pooling with genuine multi-head self-attention across | |
| # agents. (The original repo ALREADY has single-head | |
| # softmax attention via `attention=True`; this is the | |
| # real multi-head upgrade the PDF describes.) | |
| # 3. hj_residual_loss / consistency_loss -- the physics loss terms, written so | |
| # they are internally consistent (both use the geodesic | |
| # field, unlike the PDF which mixes geodesic + Euclidean). | |
| # | |
| # These modules mirror the shapes used in DQNModel.py: | |
| # input : (batch, agents, frame_history, D, W, H) for the image stream | |
| # physics: (batch, agents, C_phys, D, W, H) for the physics stream | |
| # output: (batch, agents, number_actions) | |
| # | |
| # NOTE: this file is torch code and depends on the same torch version as the | |
| # rest of the repo. It has not been executed in this environment (no GPU/torch | |
| # here), but it is structured to be a drop-in sibling of DQNModel.CommNet. | |
| import torch | |
| import torch.nn as nn | |
| # ============================================================================= | |
| # 1. Physics encoder branch (PDF Step 2) | |
| # ============================================================================= | |
| class PhysicsEncoder(nn.Module): | |
| """Encodes the (speed / potential / distance / gradient) channels into a | |
| feature vector concatenated with the image-CNN features before communication. | |
| Deliberately lightweight -- the physics maps are smooth and low-frequency.""" | |
| def __init__(self, in_channels=3, out_features=128): | |
| super().__init__() | |
| self.net = nn.Sequential( | |
| nn.Conv3d(in_channels, 16, 3, padding=1), nn.PReLU(), | |
| nn.MaxPool3d(2), | |
| nn.Conv3d(16, 32, 3, padding=1), nn.PReLU(), | |
| nn.MaxPool3d(2), | |
| nn.Conv3d(32, 32, 3, padding=0), nn.PReLU(), | |
| nn.AdaptiveAvgPool3d(1), | |
| ) | |
| self.fc = nn.Linear(32, out_features) | |
| def forward(self, x): # x: (B*agents, C, D, W, H) | |
| h = self.net(x).flatten(1) | |
| return self.fc(h) | |
| # ============================================================================= | |
| # 2. Multi-head attention communication (PDF Step 8, done properly) | |
| # ============================================================================= | |
| class MultiHeadCommNet(nn.Module): | |
| """CommNet with true multi-head self-attention between agents at each stage. | |
| Compared to the repo's CommNet: | |
| * mean pooling -> nn.MultiheadAttention (query/key/value over the agent | |
| axis), with a residual connection and LayerNorm (a transformer block). | |
| * optionally fuses a PhysicsEncoder feature into the shared representation. | |
| """ | |
| def __init__(self, agents, frame_history, number_actions, | |
| phys_channels=0, n_heads=4, xavier=True): | |
| super().__init__() | |
| self.agents = agents | |
| self.frame_history = frame_history | |
| self.use_phys = phys_channels > 0 | |
| self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| # --- shared image CNN (identical to DQNModel.CommNet backbone) -------- | |
| def conv(i, o, k, p): | |
| return nn.Conv3d(i, o, k, padding=p) | |
| self.conv0 = conv(frame_history, 32, 5, 1); self.pool0 = nn.MaxPool3d(2); self.act0 = nn.PReLU() | |
| self.conv1 = conv(32, 32, 5, 1); self.pool1 = nn.MaxPool3d(2); self.act1 = nn.PReLU() | |
| self.conv2 = conv(32, 64, 4, 1); self.pool2 = nn.MaxPool3d(2); self.act2 = nn.PReLU() | |
| self.conv3 = conv(64, 64, 3, 0); self.act3 = nn.PReLU() | |
| self.img_dim = 512 | |
| # --- physics branch --------------------------------------------------- | |
| if self.use_phys: | |
| self.phys = PhysicsEncoder(phys_channels, out_features=128) | |
| self.fuse = nn.Linear(self.img_dim + 128, self.img_dim) | |
| # --- three transformer communication blocks over the agent axis ------- | |
| self.attn = nn.ModuleList([ | |
| nn.MultiheadAttention(self.img_dim, n_heads, batch_first=True) | |
| for _ in range(3)]) | |
| self.norm = nn.ModuleList([nn.LayerNorm(self.img_dim) for _ in range(3)]) | |
| self.ff = nn.ModuleList([ | |
| nn.Sequential(nn.Linear(self.img_dim, self.img_dim), nn.PReLU()) | |
| for _ in range(3)]) | |
| # --- per-agent heads -------------------------------------------------- | |
| self.head = nn.ModuleList([ | |
| nn.Sequential(nn.Linear(self.img_dim, 256), nn.PReLU(), | |
| nn.Linear(256, number_actions)) | |
| for _ in range(agents)]) | |
| self.to(self.device) | |
| if xavier: | |
| for m in self.modules(): | |
| if isinstance(m, (nn.Conv3d, nn.Linear)): | |
| nn.init.xavier_uniform_(m.weight) | |
| def _cnn(self, x): # x: (B, frame_history, D, W, H) | |
| x = self.act0(self.pool0(self.conv0(x))) | |
| x = self.act1(self.pool1(self.conv1(x))) | |
| x = self.act2(self.pool2(self.conv2(x))) | |
| x = self.act3(self.conv3(x)) | |
| return x.reshape(x.size(0), -1)[:, :self.img_dim] | |
| def forward(self, image, physics=None): | |
| """image: (B, agents, frame_history, D, W, H) | |
| physics: (B, agents, C, D, W, H) or None""" | |
| image = image.to(self.device) / 255.0 | |
| B = image.size(0) | |
| feats = [] | |
| for i in range(self.agents): | |
| f = self._cnn(image[:, i]) | |
| if self.use_phys and physics is not None: | |
| pf = self.phys(physics[:, i].to(self.device)) | |
| f = self.fuse(torch.cat([f, pf], dim=-1)) | |
| feats.append(f) | |
| x = torch.stack(feats, dim=1) # (B, agents, img_dim) | |
| for attn, norm, ff in zip(self.attn, self.norm, self.ff): | |
| a, _ = attn(x, x, x) # attention across agents | |
| x = norm(x + a) # residual + norm | |
| x = x + ff(x) | |
| out = torch.stack([self.head[i](x[:, i]) for i in range(self.agents)], dim=1) | |
| return out.cpu() | |
| # ============================================================================= | |
| # 3. Physics loss terms (PDF Step 7, made internally consistent) | |
| # ============================================================================= | |
| def hj_residual_loss(V, F): | |
| """L_HJ = ( ||grad V|| * F - 1 )^2 on a predicted potential field V. | |
| V, F: (B, D, W, H) tensors. Uses finite differences for the gradient. | |
| Only needed if you actually *learn* V with a PINN head. If V is precomputed | |
| offline (recommended), this term is unnecessary -- V already satisfies the | |
| equation by construction. | |
| """ | |
| gx = V[:, 1:, :, :] - V[:, :-1, :, :] | |
| gy = V[:, :, 1:, :] - V[:, :, :-1, :] | |
| gz = V[:, :, :, 1:] - V[:, :, :, :-1] | |
| # pad back to equal size | |
| gx = torch.nn.functional.pad(gx, (0, 0, 0, 0, 0, 1)) | |
| gy = torch.nn.functional.pad(gy, (0, 0, 0, 1, 0, 0)) | |
| gz = torch.nn.functional.pad(gz, (0, 1, 0, 0, 0, 0)) | |
| grad_norm = torch.sqrt(gx ** 2 + gy ** 2 + gz ** 2 + 1e-8) | |
| return ((grad_norm * F - 1.0) ** 2).mean() | |
| def consistency_loss(V_pred, V_geodesic): | |
| """L_Consistency = ||V_pred - V_geodesic||^2. | |
| IMPORTANT FIX vs the PDF: the target is the GEODESIC field V_geodesic (the | |
| solution of the Eikonal equation), NOT the Euclidean distance d(x). Using | |
| Euclidean d, as the PDF does, directly contradicts L_HJ whenever the speed | |
| field F varies (which is the entire point of the physics). See EVALUATION.md | |
| for the ~48-voxel discrepancy this causes on real ADNI data. | |
| """ | |
| return ((V_pred - V_geodesic) ** 2).mean() | |
| def communication_loss(agent_features): | |
| """L_Comm = sum_i ||z_i - mean_j z_j||^2 (encourages a shared code).""" | |
| z_bar = agent_features.mean(dim=1, keepdim=True) | |
| return ((agent_features - z_bar) ** 2).mean() | |