| """Action-Conditioned Predictor for Grid-JEPA.""" |
| from typing import Optional |
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
|
|
| class DiscreteActionEmbed(nn.Module): |
| def __init__(self, num_key_actions=6, grid_size=64, action_embed_dim=64): |
| super().__init__() |
| self.num_key_actions = num_key_actions |
| self.grid_size = grid_size |
| self.num_positions = grid_size * grid_size |
| self.key_embed = nn.Embedding(num_key_actions, action_embed_dim) |
| self.pos_embed = nn.Embedding(self.num_positions, action_embed_dim) |
| def forward(self, key_action, cell_position=None): |
| emb = self.key_embed(key_action) |
| if cell_position is not None: |
| emb = emb + self.pos_embed(cell_position) |
| return emb |
|
|
| class PredictorBlock(nn.Module): |
| def __init__(self, dim, num_heads, mlp_ratio=4.0): |
| super().__init__() |
| self.norm1 = nn.LayerNorm(dim, eps=1e-6) |
| self.num_heads = num_heads |
| self.head_dim = dim // num_heads |
| self.scale = self.head_dim ** -0.5 |
| self.qkv = nn.Linear(dim, dim * 3) |
| self.proj = nn.Linear(dim, dim) |
| self.norm2 = nn.LayerNorm(dim, eps=1e-6) |
| hidden = int(dim * mlp_ratio) |
| self.mlp = nn.Sequential(nn.Linear(dim, hidden), nn.GELU(), nn.Linear(hidden, dim)) |
| def forward(self, x): |
| B, N, D = x.shape |
| xn = self.norm1(x) |
| qkv = self.qkv(xn).reshape(B, N, 3, self.num_heads, self.head_dim) |
| qkv = qkv.permute(2, 0, 3, 1, 4) |
| q, k, v = qkv[0], qkv[1], qkv[2] |
| attn = F.softmax((q @ k.transpose(-2, -1)) * self.scale, dim=-1) |
| out = (attn @ v).transpose(1, 2).reshape(B, N, D) |
| x = x + self.proj(out) |
| x = x + self.mlp(self.norm2(x)) |
| return x |
|
|
| class ActionConditionedPredictor(nn.Module): |
| def __init__(self, num_patches=4096, embed_dim=384, action_embed_dim=64, |
| predictor_depth=12, num_heads=6): |
| super().__init__() |
| self.num_patches = num_patches |
| self.embed_dim = embed_dim |
| self.mask_token = nn.Parameter(torch.zeros(1, 1, embed_dim)) |
| self.pos_embed = nn.Parameter(torch.zeros(1, num_patches, embed_dim)) |
| nn.init.trunc_normal_(self.mask_token, std=0.02) |
| nn.init.trunc_normal_(self.pos_embed, std=0.02) |
| self.action_mixer = nn.Linear(embed_dim + action_embed_dim, embed_dim) |
| self.blocks = nn.ModuleList([PredictorBlock(embed_dim, num_heads) for _ in range(predictor_depth)]) |
| self.norm = nn.LayerNorm(embed_dim, eps=1e-6) |
| def forward(self, context_repr, action_emb, target_positions): |
| B = context_repr.shape[0] |
| if target_positions.dim() == 1: |
| N_tgt = target_positions.shape[0] |
| tgt_pos = target_positions.unsqueeze(0).expand(B, -1) |
| else: |
| N_tgt = target_positions.shape[1] |
| tgt_pos = target_positions |
| tgt_pos_emb = torch.gather( |
| self.pos_embed.expand(B, -1, -1), 1, |
| tgt_pos.unsqueeze(-1).expand(-1, -1, self.embed_dim) |
| ) |
| mask_tokens = self.mask_token.expand(B, N_tgt, -1) + tgt_pos_emb |
| act_tokens = action_emb.unsqueeze(1).expand(-1, N_tgt, -1) |
| tokens = self.action_mixer(torch.cat([mask_tokens, act_tokens], dim=-1)) |
| x = torch.cat([context_repr, tokens], dim=1) |
| for block in self.blocks: |
| x = block(x) |
| x = self.norm(x) |
| return x[:, -N_tgt:, :] |
|
|
| class GridWorldPredictor(nn.Module): |
| def __init__(self, num_colors=16, embed_dim=384, predictor_depth=12, num_heads=6, |
| num_key_actions=6, grid_size=64, action_embed_dim=64): |
| super().__init__() |
| self.num_patches = grid_size * grid_size |
| self.action_embed = DiscreteActionEmbed(num_key_actions, grid_size, action_embed_dim) |
| self.predictor = ActionConditionedPredictor(self.num_patches, embed_dim, action_embed_dim, predictor_depth, num_heads) |
| self.decoder = nn.Sequential(nn.Linear(embed_dim, embed_dim), nn.GELU(), nn.Linear(embed_dim, num_colors)) |
| def predict_next_state(self, current_repr, action_key, action_pos=None): |
| B = current_repr.shape[0] |
| action_emb = self.action_embed(action_key, action_pos) |
| all_pos = torch.arange(self.num_patches, device=current_repr.device).unsqueeze(0).expand(B, -1) |
| return self.predictor(current_repr, action_emb, all_pos) |
| def decode_to_grid(self, repr): |
| return self.decoder(repr) |
|
|
| if __name__ == "__main__": |
| B, N_ctx, D, N_tgt = 2, 60, 128, 40 |
| context = torch.randn(B, N_ctx, D) |
| action_key = torch.randint(0, 6, (B,)) |
| action_pos = torch.randint(0, 100, (B,)) |
| target_pos = torch.randint(0, 100, (N_tgt,)) |
| embed = DiscreteActionEmbed(6, 10, 64) |
| print(f"Action emb: {embed(action_key, action_pos).shape}") |
| pred = ActionConditionedPredictor(100, D, 64, 4, 4) |
| print(f"Pred: {pred(context, embed(action_key, action_pos), target_pos).shape}") |
| wp = GridWorldPredictor(10, D, 4, 4, 6, 10, 64) |
| print(f"Next state: {wp.predict_next_state(context, action_key, action_pos).shape}") |
| print(f"Decode: {wp.decode_to_grid(wp.predict_next_state(context, action_key, action_pos)).shape}") |
|
|
