initial commit

src/velocity_prediction/model.py

@@ -0,0 +1,170 @@
+"""
+VelocityPredictionModel: CNN + PoseMLP → concat → GRU → Head → [vx_body, vy_body].
+
+Architecture:
+    Event frame (1, H, W) ──► CNN ──┐
+    Tilt angles    (3,)    ──► MLP ──┤──► concat ──► GRU ──► Head ──► [vx, vy]
+"""
+
+import torch
+import torch.nn as nn
+
+from src.velocity_prediction.config import model_cfg
+
+
+class CNNEncoder(nn.Module):
+    """
+    4-layer ConvNet with BatchNorm, ReLU, MaxPool, ending with Global Avg Pool.
+
+    Input:  (B, S, 1, H, W)  — processed per-frame (flattened to (B*S, 1, H, W))
+    Output: (B, S, C_out)     — per-frame feature vectors
+    """
+
+    def __init__(self, cfg=model_cfg.cnn):
+        super().__init__()
+        channels = cfg.channels
+        in_ch = cfg.in_channels
+
+        layers = []
+        for out_ch in channels:
+            layers.extend([
+                nn.Conv2d(in_ch, out_ch, kernel_size=cfg.kernel_size, padding=cfg.kernel_size // 2),
+                # nn.BatchNorm2d(out_ch) if cfg.use_bn else nn.Identity(),
+                nn.Identity(),
+                nn.LeakyReLU(inplace=True),
+                nn.MaxPool2d(cfg.pool_size),
+            ])
+            in_ch = out_ch
+
+        self.conv = nn.Sequential(*layers)
+        self.gap = nn.AdaptiveAvgPool2d(1)
+        self.out_dim = channels[-1]
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            x: (B, S, 1, H, W)  event frame sequence
+        Returns:
+            features: (B, S, C_out)
+        """
+        B, S, C, H, W = x.shape
+        x = x.view(B * S, C, H, W)          # (B*S, 1, H, W)
+        x = self.conv(x)                     # (B*S, C_out, H', W')
+        x = self.gap(x)                      # (B*S, C_out, 1, 1)
+        x = x.view(B, S, self.out_dim)       # (B, S, C_out)
+        return x
+
+
+class PoseMLP(nn.Module):
+    """
+    Encode tilt rotation vector (3,) into a compact feature vector.
+
+    Input:  (B, S, 3)
+    Output: (B, S, output_dim)
+    """
+
+    def __init__(self, cfg=model_cfg.pose_mlp):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Linear(cfg.input_dim, cfg.hidden_dim),
+            nn.LeakyReLU(inplace=True),
+            nn.Linear(cfg.hidden_dim, cfg.output_dim),
+            nn.LeakyReLU(inplace=True),
+        )
+        self.out_dim = cfg.output_dim
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        x: (B, S, 3)  →  (B, S, output_dim)
+        """
+        B, S, D = x.shape
+        x = x.view(B * S, D)
+        x = self.net(x)
+        x = x.view(B, S, self.out_dim)
+        return x
+
+
+class VelocityPredictionModel(nn.Module):
+    """
+    Full model: CNN + PoseMLP → concat → GRU → Head → [vx, vy].
+
+    Input:
+        events: (B, S, 1, H, W)
+        tilt:   (B, S, 3)
+    Output:
+        v_body: (B, 2)  — body-frame [vx, vy] for the last frame in the sequence
+    """
+
+    def __init__(self, cnn_cfg=model_cfg.cnn, pose_cfg=model_cfg.pose_mlp,
+                 gru_cfg=model_cfg.gru, head_cfg=model_cfg.head):
+        super().__init__()
+
+        self.cnn = CNNEncoder(cnn_cfg)
+        self.pose_mlp = PoseMLP(pose_cfg)
+
+        fused_dim = self.cnn.out_dim + self.pose_mlp.out_dim  # 256 + 64 = 320
+
+        self.gru = nn.GRU(
+            input_size=fused_dim,
+            hidden_size=gru_cfg.hidden_size,
+            num_layers=gru_cfg.num_layers,
+            dropout=gru_cfg.dropout if gru_cfg.num_layers > 1 else 0.0,
+            batch_first=True,
+        )
+
+        self.head = nn.Sequential(
+            nn.Linear(gru_cfg.hidden_size, head_cfg.hidden_dim),
+            nn.LeakyReLU(inplace=True),
+            nn.Linear(head_cfg.hidden_dim, head_cfg.output_dim),
+        )
+
+        # # Small init for the final layer: start from near-zero output
+        # self.head[-1].weight.data.mul_(0.01)
+        # self.head[-1].bias.data.zero_()
+
+    def forward(self, events: torch.Tensor, tilt: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            events: (B, S, 1, H, W)
+            tilt:   (B, S, 3)
+
+        Returns:
+            v_body: (B, 2)  predicted body-frame [vx, vy] at the last timestep
+        """
+        # Per-frame encoding
+        cnn_feat = self.cnn(events)           # (B, S, 256)
+        pose_feat = self.pose_mlp(tilt)       # (B, S, 64)
+
+        # Fuse per frame
+        fused = torch.cat([cnn_feat, pose_feat], dim=-1)  # (B, S, 320)
+
+        # GRU temporal modelling
+        gru_out, h_n = self.gru(fused)        # gru_out: (B, S, 128), h_n: (1, B, 128)
+
+        # Use last hidden state
+        last_hidden = h_n[-1]                 # (B, 128)
+
+        # Head regression
+        v_body = self.head(last_hidden)       # (B, 2)
+        return v_body
+
+
+def count_parameters(model: nn.Module) -> int:
+    """Count trainable parameters."""
+    return sum(p.numel() for p in model.parameters() if p.requires_grad)
+
+
+if __name__ == "__main__":
+    # Quick sanity check
+    model = VelocityPredictionModel()
+    total = count_parameters(model)
+    print(f"Total trainable parameters: {total:,} ({total/1e6:.3f} M)")
+
+    # Forward pass test
+    B, S, H, W = 4, 8, 240, 320
+    events = torch.randn(B, S, 1, H, W)
+    tilt = torch.randn(B, S, 3)
+    out = model(events, tilt)
+    print(f"Input events: {events.shape}")
+    print(f"Input tilt:   {tilt.shape}")
+    print(f"Output:       {out.shape}  (should be [4, 2])")