| """Complete Grid-JEPA system.""" |
| import random |
| import torch |
| import torch.nn as nn |
| import torch.optim as optim |
| from torch.utils.data import DataLoader |
|
|
| try: |
| from .encoder import GridPatchEmbed, ViTEncoder |
| from .predictor import DiscreteActionEmbed, ActionConditionedPredictor |
| except ImportError: |
| from encoder import GridPatchEmbed, ViTEncoder |
| from predictor import DiscreteActionEmbed, ActionConditionedPredictor |
|
|
| class GridJEPA(nn.Module): |
| def __init__(self, num_colors=16, embed_dim=384, encoder_depth=12, predictor_depth=12, |
| num_heads=6, mlp_ratio=4.0, max_grid_size=64, ema_decay=0.996, |
| num_key_actions=6, action_embed_dim=64, use_action_conditioning=True): |
| super().__init__() |
| self.num_colors = num_colors |
| self.embed_dim = embed_dim |
| self.max_grid_size = max_grid_size |
| self.num_patches = max_grid_size * max_grid_size |
| self.use_action_conditioning = use_action_conditioning |
| self.ema_decay = ema_decay |
|
|
| self.patch_embed = GridPatchEmbed(num_colors, embed_dim, max_grid_size) |
| self.context_encoder = ViTEncoder(embed_dim, encoder_depth, num_heads, mlp_ratio) |
| self.target_encoder = ViTEncoder(embed_dim, encoder_depth, num_heads, mlp_ratio) |
| self.target_encoder.load_state_dict(self.context_encoder.state_dict()) |
| for p in self.target_encoder.parameters(): |
| p.requires_grad = False |
|
|
| self.action_embed = DiscreteActionEmbed(num_key_actions, max_grid_size, action_embed_dim) |
| self.predictor = ActionConditionedPredictor(self.num_patches, embed_dim, action_embed_dim, predictor_depth, num_heads) |
|
|
| def forward(self, grid, context_mask, target_mask, action_key=None, action_pos=None): |
| B = grid.shape[0] |
| N = self.num_patches |
| x = self.patch_embed(grid) |
| ctx_x = x.clone() |
| ctx_x[~context_mask] = 0 |
| context_repr = self.context_encoder(ctx_x) |
| with torch.no_grad(): |
| target_repr = self.target_encoder(x) |
|
|
| target_indices = [target_mask[b].nonzero(as_tuple=True)[0] for b in range(B)] |
|
|
| if self.use_action_conditioning and action_key is not None: |
| action_emb = self.action_embed(action_key, action_pos) |
| else: |
| action_emb = torch.zeros(B, 64, device=grid.device) |
|
|
| losses = [] |
| for b in range(B): |
| tgt_idx = target_indices[b] |
| if len(tgt_idx) == 0: |
| continue |
| pred = self.predictor(context_repr[b:b+1], action_emb[b:b+1], tgt_idx.unsqueeze(0)) |
| tgt = target_repr[b, tgt_idx].unsqueeze(0) |
| losses.append((pred - tgt).pow(2).sum(dim=-1).mean()) |
|
|
| if len(losses) == 0: |
| return torch.tensor(0.0, device=grid.device), {} |
|
|
| total = torch.stack(losses).mean() |
| metrics = {"loss": total.item(), "num_targets": sum(len(t) for t in target_indices) / B} |
| return total, metrics |
|
|
| def update_ema(self): |
| with torch.no_grad(): |
| for pt, pc in zip(self.target_encoder.parameters(), self.context_encoder.parameters()): |
| pt.data.mul_(self.ema_decay).add_(pc.data, alpha=1 - self.ema_decay) |
|
|
| def encode(self, grid): |
| return self.context_encoder(self.patch_embed(grid)) |
|
|
| @torch.no_grad() |
| def encode_target(self, grid): |
| return self.target_encoder(self.patch_embed(grid)) |
|
|
| if __name__ == "__main__": |
| B, G, C = 2, 10, 10 |
| device = "cuda" if torch.cuda.is_available() else "cpu" |
| model = GridJEPA(C, 128, 4, 4, 4, max_grid_size=G).to(device) |
| grid = torch.randint(0, C, (B, G, G), device=device) |
| N = G * G |
| ctx_mask = torch.zeros(B, N, dtype=torch.bool, device=device) |
| tgt_mask = torch.zeros(B, N, dtype=torch.bool, device=device) |
| for b in range(B): |
| idx = list(range(N)) |
| random.shuffle(idx) |
| split = N // 2 |
| ctx_mask[b, idx[:split]] = True |
| tgt_mask[b, idx[split:]] = True |
| action_key = torch.randint(0, 6, (B,), device=device) |
| action_pos = torch.randint(0, N, (B,), device=device) |
| loss, metrics = model(grid, ctx_mask, tgt_mask, action_key, action_pos) |
| print(f"Loss: {loss.item():.2f}, Metrics: {metrics}") |
| model.update_ema() |
| print("EMA OK") |
| print(f"Encode: {model.encode(grid).shape}") |
|
|
