Upload lewm_model.py

lewm_model.py

@@ -0,0 +1,211 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+class SIGReg(nn.Module):
+    """Sketched Isotropic Gaussian Regularizer (SIGReg)
+    Encourages latent embeddings to match an isotropic Gaussian N(0, I).
+    Uses random 1D projections + Epps-Pulley test statistic.
+    Input shape: (N, D) or (T, B, D) — any shape where last dim is D.
+    """
+    def __init__(self, knots=17, num_proj=1024, t_min=0.2, t_max=4.0):
+        super().__init__()
+        self.num_proj = num_proj
+        t = torch.linspace(t_min, t_max, knots, dtype=torch.float32)
+        dt = (t_max - t_min) / (knots - 1)
+        weights = torch.full((knots,), 2 * dt, dtype=torch.float32)
+        weights[[0, -1]] = dt
+        window = torch.exp(-t.square() / 2.0)
+        self.register_buffer("t", t)
+        self.register_buffer("phi", window)
+        self.register_buffer("weights", weights * window)
+
+    def forward(self, proj):
+        # proj: (*, D) where D is embedding dim
+        D = proj.size(-1)
+        A = torch.randn(D, self.num_proj, device=proj.device, dtype=proj.dtype)
+        A = A.div_(A.norm(p=2, dim=0, keepdim=True))
+        x = proj @ A  # (*, num_proj)
+        x_t = x.unsqueeze(-1) * self.t  # (*, num_proj, knots)
+        cos_mean = x_t.cos().mean(dim=-3)  # (num_proj, knots)
+        sin_mean = x_t.sin().mean(dim=-3)  # (num_proj, knots)
+        err = (cos_mean - self.phi).square() + sin_mean.square()
+        statistic = (err @ self.weights) * proj.size(-2)
+        return statistic.mean()
+
+class AdaLNBlock(nn.Module):
+    """Transformer block with Adaptive Layer Normalization (AdaLN-zero) conditioning."""
+    def __init__(self, dim, heads=16, mlp_ratio=4.0, dropout=0.1):
+        super().__init__()
+        self.norm1 = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
+        self.norm2 = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
+        self.attn = nn.MultiheadAttention(dim, heads, dropout=dropout, batch_first=True)
+        mlp_dim = int(dim * mlp_ratio)
+        self.mlp = nn.Sequential(
+            nn.Linear(dim, mlp_dim),
+            nn.GELU(),
+            nn.Dropout(dropout),
+            nn.Linear(mlp_dim, dim),
+            nn.Dropout(dropout),
+        )
+        self.adaLN_modulation = nn.Sequential(
+            nn.SiLU(),
+            nn.Linear(dim, 6 * dim, bias=True)
+        )
+        nn.init.constant_(self.adaLN_modulation[-1].weight, 0.0)
+        nn.init.constant_(self.adaLN_modulation[-1].bias, 0.0)
+
+    def modulate(self, x, shift, scale):
+        return x * (1 + scale) + shift
+
+    def forward(self, x, c):
+        if c.size(1) != x.size(1):
+            pad_len = x.size(1) - c.size(1)
+            c = F.pad(c, (0, 0, 0, pad_len), mode='constant', value=0)
+        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = (
+            self.adaLN_modulation(c).chunk(6, dim=-1)
+        )
+        normed = self.modulate(self.norm1(x), shift_msa, scale_msa)
+        attn_out, _ = self.attn(normed, normed, normed, need_weights=False)
+        x = x + gate_msa * attn_out
+        normed = self.modulate(self.norm2(x), shift_mlp, scale_mlp)
+        x = x + gate_mlp * self.mlp(normed)
+        return x
+
+class Predictor(nn.Module):
+    """Causal Transformer predictor with AdaLN action conditioning."""
+    def __init__(self, dim, depth=6, heads=16, mlp_ratio=4.0, dropout=0.1, max_len=512):
+        super().__init__()
+        self.pos_emb = nn.Parameter(torch.randn(1, max_len, dim) * 0.02)
+        self.blocks = nn.ModuleList([
+            AdaLNBlock(dim, heads, mlp_ratio, dropout) for _ in range(depth)
+        ])
+        self.norm = nn.LayerNorm(dim, eps=1e-6)
+
+    def forward(self, x, act_emb):
+        B, T, D = x.shape
+        x = x + self.pos_emb[:, :T, :]
+        for block in self.blocks:
+            x = block(x, act_emb)
+        x = self.norm(x)
+        return x
+
+class ActionEncoder(nn.Module):
+    """Encode continuous actions into embedding space."""
+    def __init__(self, action_dim, out_dim, hidden_dim=256):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Linear(action_dim, hidden_dim),
+            nn.ReLU(),
+            nn.Linear(hidden_dim, hidden_dim),
+            nn.ReLU(),
+            nn.Linear(hidden_dim, out_dim),
+        )
+
+    def forward(self, a):
+        return self.net(a)
+
+class LeWorldModel(nn.Module):
+    """
+    LeWorldModel (LeWM): Stable end-to-end JEPA world model from pixels.
+    Encoder: ViT-Tiny + BatchNorm projector.
+    Predictor: Causal Transformer with AdaLN-zero action conditioning.
+    """
+    def __init__(
+        self,
+        img_size=224,
+        patch_size=14,
+        in_chans=3,
+        embed_dim=192,
+        encoder_depth=12,
+        encoder_heads=3,
+        predictor_depth=6,
+        predictor_heads=16,
+        action_dim=2,
+        latent_dim=192,
+        dropout=0.1,
+        max_len=512,
+    ):
+        super().__init__()
+        self.patch_embed = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)
+        num_patches = (img_size // patch_size) ** 2
+        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
+        self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, embed_dim))
+        self.pos_drop = nn.Dropout(p=dropout)
+        encoder_layer = nn.TransformerEncoderLayer(
+            d_model=embed_dim,
+            nhead=encoder_heads,
+            dim_feedforward=int(embed_dim * 4),
+            dropout=dropout,
+            activation='gelu',
+            batch_first=True,
+            norm_first=True,
+        )
+        self.encoder_transformer = nn.TransformerEncoder(encoder_layer, num_layers=encoder_depth)
+        self.encoder_norm = nn.LayerNorm(embed_dim, eps=1e-6)
+
+        self.projector = nn.Sequential(
+            nn.Linear(embed_dim, latent_dim),
+            nn.BatchNorm1d(latent_dim),
+        )
+
+        self.action_encoder = ActionEncoder(action_dim, latent_dim)
+
+        self.predictor = Predictor(
+            dim=latent_dim,
+            depth=predictor_depth,
+            heads=predictor_heads,
+            dropout=dropout,
+            max_len=max_len,
+        )
+
+        self.pred_projector = nn.Sequential(
+            nn.Linear(latent_dim, latent_dim),
+            nn.BatchNorm1d(latent_dim),
+        )
+
+        self._init_weights()
+
+    def _init_weights(self):
+        nn.init.normal_(self.cls_token, std=0.02)
+        nn.init.normal_(self.pos_embed, std=0.02)
+
+    def encode(self, obs):
+        if obs.dim() == 5:
+            B, T, C, H, W = obs.shape
+            obs = obs.reshape(B * T, C, H, W)
+            reshape_back = True
+        else:
+            B = obs.size(0)
+            reshape_back = False
+
+        x = self.patch_embed(obs)
+        x = x.flatten(2).transpose(1, 2)
+        cls_tokens = self.cls_token.expand(x.size(0), -1, -1)
+        x = torch.cat([cls_tokens, x], dim=1)
+        x = x + self.pos_embed
+        x = self.pos_drop(x)
+        x = self.encoder_transformer(x)
+        x = self.encoder_norm(x)
+        cls = x[:, 0]
+        z = self.projector(cls)
+
+        if reshape_back:
+            z = z.reshape(B, T, -1)
+        return z
+
+    def predict(self, emb, actions):
+        B, T, D = emb.shape
+        act_emb = self.action_encoder(actions)
+        pad = torch.zeros(B, 1, D, device=act_emb.device, dtype=act_emb.dtype)
+        act_emb = torch.cat([act_emb, pad], dim=1)
+        preds = self.predictor(emb, act_emb)
+        preds = preds.reshape(B * T, D)
+        preds = self.pred_projector(preds)
+        preds = preds.reshape(B, T, D)
+        return preds
+
+    def forward(self, obs, actions):
+        emb = self.encode(obs)
+        pred_emb = self.predict(emb, actions)
+        return emb, pred_emb