Spaces:
Running on Zero
Running on Zero
File size: 10,965 Bytes
dcc4594 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 | from models.utils import ModelType
import torch
import torch.nn as nn
import math
from models.positional_encoding import build_spatial_only_pe
from models.st_transformer import STTransformer
from einops import repeat
class DynamicsModel(nn.Module):
def __init__(self, frame_size=(128, 128), patch_size=4, embed_dim=128, num_heads=8,
hidden_dim=128, num_blocks=4, num_bins=4, n_actions=8, conditioning_dim=3, latent_dim=5,
use_moe=False, num_experts=4, top_k_experts=2, moe_aux_loss_coeff=0.01):
super().__init__()
H, W = frame_size
codebook_size = num_bins**latent_dim
self.latent_embed = nn.Linear(latent_dim, embed_dim)
self.transformer = STTransformer(
embed_dim, num_heads, hidden_dim, num_blocks, causal=True,
conditioning_dim=conditioning_dim,
use_moe=use_moe, num_experts=num_experts,
top_k_experts=top_k_experts, moe_aux_loss_coeff=moe_aux_loss_coeff,
)
self.output_mlp = nn.Linear(embed_dim, codebook_size)
# shared spatial-only PE (zeros in temporal tail)
pe_spatial = build_spatial_only_pe((H, W), patch_size, embed_dim, device='cpu', dtype=torch.float32) # [1,P,E]
self.register_buffer("pos_spatial_dec", pe_spatial, persistent=False)
# learnable mask token latent
# TODO; try leanable mask embedding in embed space instead of latent space
self.mask_token = nn.Parameter(torch.randn(1, 1, 1, latent_dim) * 0.02) # [1, 1, 1, L]
def forward(self, discrete_latents, training=True, conditioning=None, targets=None):
# discrete_latents: [B, T, P, L]
# targets: [B, T, P] indices
# conditioning: [B, T, A]
B, T, P, L = discrete_latents.shape
# convert latents to float for embedding
discrete_latents = discrete_latents.to(dtype=torch.float32)
# apply MaskGIT random masking during training
if training and self.training:
# per-batch mask ratio in [0.5, 1.0)
mask_ratio = 0.5 + torch.rand((), device=discrete_latents.device) * 0.5
mask_positions = (torch.rand(B, T, P, device=discrete_latents.device) < mask_ratio) # [B, T, P]
# guarantee at least one unmasked temporal anchor per (B, P)
# pick a random timestep for each (B,P) and force it to unmask
anchor_idx = torch.randint(0, T, (B, P), device=discrete_latents.device) # [B, P]
mask_positions[torch.arange(B)[:, None], anchor_idx, torch.arange(P)[None, :]] = False # [B, T, P]
# replace selected latents with mask tokens
mask_token = repeat(self.mask_token.to(discrete_latents.device, discrete_latents.dtype), '1 1 1 L -> B T P L', B=B, T=T, P=P) # [B, T, P, L]
discrete_latents = torch.where(mask_positions.unsqueeze(-1), mask_token, discrete_latents) # [B, T, P, L]
else:
mask_positions = None
embeddings = self.latent_embed(discrete_latents) # [B, T, P, E]
# add spatial PE (affects only first 2/3 of dimensions)
# STTransformer adds temporal PE to last 1/3 of dimensions
embeddings = embeddings + self.pos_spatial_dec.to(embeddings.device, embeddings.dtype)
transformed = self.transformer(embeddings, conditioning=conditioning) # [B, T, P, E]
# transform to logits for each token in codebook
predicted_logits = self.output_mlp(transformed) # [B, T, P, L^D]
# compute masked cross-entropy loss
loss = None
if training and self.training:
assert targets is not None, "target indices are needed for training"
Ld = predicted_logits.shape[-1] # L^D
logits_flat = predicted_logits.reshape(-1, Ld) # [(B*T*P), L^D]
targets_flat = targets.reshape(-1) # [(B*T*P)]
mask_flat = mask_positions.reshape(-1).to(torch.float32) # [(B*T*P)]
loss_per = nn.functional.cross_entropy(logits_flat, targets_flat, reduction='none') # [(B*T*P)]
denom = mask_flat.sum().clamp_min(1.0)
loss = (loss_per * mask_flat).sum() / denom
return predicted_logits, mask_positions, loss # logits, mask, optional loss
def exp_schedule_torch(self, t, T, P_total, k, device):
# t: current step, T: total steps, P_total: total masked positions across the horizon window
# exp schedule is P_total * (1 - exp(k * t / T)) / (1 - exp(k))
x = t / max(T, 1)
k_tensor = torch.tensor(k, device=device)
result = P_total * torch.expm1(k_tensor * x) / torch.expm1(k_tensor)
if t == T - 1:
return torch.tensor(P_total, dtype=result.dtype, device=device)
return result
@torch.no_grad()
def forward_inference(self, context_latents, prediction_horizon, num_steps, index_to_latents_fn, conditioning=None, schedule_k=5.0, temperature: float = 0.0):
# MaskGIT-style iterative decoding across all prediction horizon steps
# context_latents: [B, T_ctx, P, L]
# T_ctx=context timesteps, H=prediction horizon, K=codebook size
device = context_latents.device
dtype = context_latents.dtype
B, T_ctx, P, L = context_latents.shape # B, T_ctx, P, L
H = int(prediction_horizon) # number of horizon steps to decode
# append prediction_horizon masked frame latents to predict dynamics on
mask_latents = self.mask_token.to(device, dtype).expand(B, H, P, -1) # [B, H, P, L]
input_latents = torch.cat([context_latents, mask_latents], dim=1) # [B, T_ctx+H, P, L]
mask = torch.ones(B, H, P, 1, dtype=torch.bool, device=device) # [B, H, P, 1]
P_total = H * P # total masked positions across the horizon window
for m in range(num_steps):
n_tokens_raw = self.exp_schedule_torch(m, num_steps, P_total, schedule_k, device)
# predict logits for current input
logits, _, _ = self.forward(input_latents, training=False, conditioning=conditioning, targets=None) # [B, T_ctx+H, P, L^D]
# temperature scaling
if temperature and temperature > 0:
scaled_logits = logits / float(temperature)
else:
scaled_logits = logits
probs = torch.softmax(scaled_logits, dim=-1) # [B, T_ctx+H, P, L^D]
# confidence for unmask selection always from max probability
max_probs, _ = torch.max(probs, dim=-1) # [B, T_ctx+H, P]
# choose indices either via argmax (temperature==0) or sampling
if temperature and temperature > 0:
Bc, Tc, Pc, Ld = probs.shape # Bc=B, Tc=T_ctx+H, Pc=P, Ld=L^D
sampled = torch.distributions.Categorical(probs=probs.reshape(-1, Ld)).sample()
predicted_indices = sampled.view(Bc, Tc, Pc) # [B, T_ctx+H, P]
else:
_, predicted_indices = torch.max(probs, dim=-1) # [B, T_ctx+H, P]
horizon_probs = max_probs[:, -H:, :] # [B, H, P]
# for each batch element, select tokens to unmask from all masked positions
for b in range(B):
masked_mask_all = mask[b, :, :, 0] # [H, P]
masked_flat = masked_mask_all.view(-1) # [H*P]
masked_flat_idx = torch.where(masked_flat)[0] # [num_masked]
if masked_flat_idx.numel() == 0:
continue
# TODO: try to clean this
num_masked_b = int(masked_flat_idx.numel())
prev_b = P_total - num_masked_b
target_unmasked = int(torch.ceil(n_tokens_raw).item())
k_floor = max(P_total // 16, 1)
k_b = max(k_floor, min(max(target_unmasked - prev_b, 0), num_masked_b))
pos_probs_flat = horizon_probs[b].contiguous().view(-1)[masked_flat_idx] # [num_masked]
if pos_probs_flat.numel() > k_b:
top_idx = torch.topk(pos_probs_flat, k_b, largest=True).indices
sel_flat = masked_flat_idx[top_idx]
else:
sel_flat = masked_flat_idx
# map back to (h, p)
h_sel = torch.div(sel_flat, P, rounding_mode='floor') # [k_b]
p_sel = sel_flat % P # [k_b]
# group by unique h and write sampled tokens to input tensor
if h_sel.numel() > 0:
unique_h = torch.unique(h_sel, sorted=True)
for uh in unique_h:
mask_h = (h_sel == uh)
p_list = p_sel[mask_h]
if p_list.numel() == 0:
continue
t_abs = T_ctx + int(uh.item()) # absolute time index in [T_ctx, T_ctx+H-1] (for current horizon step)
idx_sel = predicted_indices[b:b+1, t_abs:t_abs+1, p_list] # [1,1,P_sel]
pred_latents_sel = index_to_latents_fn(idx_sel) # [1,1,P_sel,L]
input_latents[b:b+1, t_abs:t_abs+1, p_list] = pred_latents_sel
mask[b, int(uh.item()), p_list, 0] = False
# early exit if all horizon tokens are unmasked
if not mask[:, :, :, 0].any(): # mask: [B,H,P,1]
break
# final completion: fill any remaining masked tokens across all horizon steps via argmax
# TODO: try removing
if mask[:, :, :, 0].any():
logits, _, _ = self.forward(input_latents, training=False, conditioning=conditioning, targets=None) # [B, T_ctx+H, P, L^D]
if temperature and temperature > 0:
scaled_logits = logits / float(temperature)
else:
scaled_logits = logits
probs = torch.softmax(scaled_logits, dim=-1) # [B, T_ctx+H, P, L^D]
_, predicted_indices = torch.max(probs, dim=-1) # [B, T_ctx+H, P]
for b in range(B):
h_idx, p_idx = torch.where(mask[b, :, :, 0]) # both [N_remaining]
if h_idx.numel() == 0:
continue
unique_h = torch.unique(h_idx, sorted=True)
for uh in unique_h:
mask_h = (h_idx == uh)
p_list = p_idx[mask_h]
if p_list.numel() == 0:
continue
t_abs = T_ctx + int(uh.item()) # absolute time index
idx_sel = predicted_indices[b:b+1, t_abs:t_abs+1, p_list] # [1,1,P_sel]
pred_latents_sel = index_to_latents_fn(idx_sel) # [1,1,P_sel,L]
input_latents[b:b+1, t_abs:t_abs+1, p_list] = pred_latents_sel
mask[b, int(uh.item()), p_list, 0] = False
return input_latents # [B, T_ctx + H, P, L]
@property
def model_type(self) -> str:
return ModelType.DynamicsModel |