Class DynamicsModel

models/dynamics.py:9–207

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]