diff --git a/latent_action_model/core/lam_lightinng.py b/latent_action_model/core/lam_lightinng.py index cf1b883..347a691 100644 --- a/latent_action_model/core/lam_lightinng.py +++ b/latent_action_model/core/lam_lightinng.py @@ -898,6 +898,34 @@ class VJEPA_LAM(LightningModule): plt.savefig(f"{filename}.png", bbox_inches="tight", pad_inches=0.0) plt.close() + def on_save_checkpoint(self, checkpoint: Dict[str, Any]) -> None: + """Drop the frozen visual encoder from the checkpoint. + + `lam.vision_encoder.*` is never trained and is re-loaded from + `vision_model_id` on every `LatentLAMModel.__init__`, so persisting it + only wastes disk. With VGGT-1B that is 909 M params -> ~3.6 GB per + checkpoint file. + """ + state_dict = checkpoint.get("state_dict") + if not state_dict: + return + for key in [k for k in state_dict if k.startswith("lam.vision_encoder.")]: + del state_dict[key] + + def on_load_checkpoint(self, checkpoint: Dict[str, Any]) -> None: + """Re-inject the frozen encoder weights stripped by `on_save_checkpoint`. + + They are already loaded in `__init__`, so copying them back from the live + module lets Lightning keep loading strictly instead of silencing genuinely + missing trained weights. + """ + state_dict = checkpoint.get("state_dict") + if state_dict is None: + return + for key, value in self.state_dict().items(): + if key.startswith("lam.vision_encoder.") and key not in state_dict: + state_dict[key] = value + def configure_optimizers(self) -> Any: if self.exclude_bias_norm_from_wd: param_groups, decay_names, no_decay_names = self._build_optimizer_param_groups() diff --git a/latent_action_model/core/lam_model.py b/latent_action_model/core/lam_model.py index c06dbd9..8bb4f0b 100644 --- a/latent_action_model/core/lam_model.py +++ b/latent_action_model/core/lam_model.py @@ -596,6 +596,10 @@ def load_latent_action_model(ckpt_path, yaml_path): ckpt_keys = set(new_ckpt.keys()) missing_keys = sorted(list(model_keys - ckpt_keys)) + # `VJEPA_LAM.on_save_checkpoint` strips the frozen visual encoder, which + # `LatentLAMModel.__init__` has already rebuilt from `vision_model_id`. + # Its absence is expected, not an error. + missing_keys = [k for k in missing_keys if not k.startswith("vision_encoder.")] unexpected_keys = sorted(list(ckpt_keys - model_keys)) shape_mismatches = [] for k in sorted(model_keys & ckpt_keys): @@ -615,6 +619,12 @@ def load_latent_action_model(ckpt_path, yaml_path): error_lines += [f" - {k}: model{ms} vs checkpoint{cs}" for k, ms, cs in shape_mismatches] raise RuntimeError("\n".join(error_lines)) + # Keep strict=True meaningful for the trained weights by filling the stripped + # encoder entries from the freshly built (already pretrained) module. + for key, value in model_state.items(): + if key.startswith("vision_encoder.") and key not in new_ckpt: + new_ckpt[key] = value + latent_action_model.load_state_dict(new_ckpt, strict=True) for p in latent_action_model.parameters(): p.requires_grad = False diff --git a/latent_action_model/core/vjepa_encoder.py b/latent_action_model/core/vjepa_encoder.py index 5d5eebe..e60d7ca 100644 --- a/latent_action_model/core/vjepa_encoder.py +++ b/latent_action_model/core/vjepa_encoder.py @@ -391,6 +391,244 @@ class DINOv3Encoder(nn.Module): return features[0].detach() if isinstance(n, int) else [f.detach() for f in features] + +# VGGT-1B aggregator output dim = 2 * embed_dim (frame-attention ‖ global-attention concat). +VGGT_FEATURE_DIM = 2048 +VGGT_DEFAULT_INPUT_SIZE = 518 + + +class VGGTEncoder(nn.Module): + """Frozen VGGT-1B aggregator, drop-in compatible with :class:`DINOv3Encoder`. + + VGGT is a 3D-geometry backbone rather than a 2D semantic one, so four + adaptations are needed to keep the rest of the LAM stack byte-for-byte + unchanged: + + 1. ``video_aug`` hands us ImageNet-normalized pixels, but the VGGT aggregator + applies its *own* ImageNet normalization internally and expects [0, 1]. + We de-normalize before the forward pass. + 2. VGGT is patch-14 at 518px -> a 37x37 grid, while the LAM decoder is wired + for ``LAM_PATCH_SIZE=16`` -> 16x16. We average-pool 37x37 down to 16x16 so + the token count matches and no downstream shape changes. + 3. VGGT's global-attention blocks mix *all* frames stacked on the S axis, and + ``LatentLAMModel._run`` concatenates (enc_t, enc_T, dec_t, dec_T) into a + single ``encode`` call. Encoding with S>1 would leak the future frame into + the current one, so we force S=1 and put every frame on the batch axis. + That makes frames independent exactly like DINOv3 (verified bit-identical + against per-frame calls; S-stacking differs by max|delta|~20). + 4. ``Aggregator.forward`` returns a length-``depth`` list where uncached + layers are ``None`` (only blocks {4, 11, 17, 23} are kept). Indexing it + with the shipped ``latent_layer_to_use=-2`` would yield ``None``, so we + index into the compacted list of cached layers instead: -1 -> block 23, + -2 -> block 17. + """ + + def __init__( + self, + model_id: str = "facebook/VGGT-1B", + num_latent_layers: int = 1, + norm_layer_type: str = "l2", + enable_norm: bool = False, + image_size: int = 256, + target_grid: int = 16, + vggt_input_size: int = VGGT_DEFAULT_INPUT_SIZE, + ): + super().__init__() + from vggt.models.vggt import VGGT + + self.device = torch.device("cpu") + self.model_id = model_id + self.num_latent_layers = max(int(num_latent_layers), 1) + self.norm_layer_type = norm_layer_type + self.enable_norm = enable_norm + self.feature_dim = VGGT_FEATURE_DIM + # `image_size` is what the LAM pipeline feeds us (256); `vggt_input_size` + # is what we upsample to internally. Reporting 256 keeps LatentLAMModel's + # image_hw consistency check quiet. + self.image_size = int(image_size) + self.vggt_input_size = int(vggt_input_size) + self.target_grid = int(target_grid) + self.patch_size = self.image_size // self.target_grid + + # Prediction heads are dead weight here: skip building them entirely + # (~500M params) and let strict=False drop their checkpoint entries. + model = VGGT( + enable_camera=False, + enable_point=False, + enable_depth=False, + enable_track=False, + ) + state_dict = _load_vggt_state_dict(model_id) + missing, unexpected = model.load_state_dict(state_dict, strict=False) + aggregator_missing = [k for k in missing if k.startswith("aggregator.")] + if aggregator_missing: + raise RuntimeError( + f"VGGT aggregator weights are incomplete: {len(aggregator_missing)} missing keys, " + f"e.g. {aggregator_missing[:5]}" + ) + print( + f"[VGGTEncoder] loaded `{model_id}` | missing={len(missing)} unexpected={len(unexpected)} " + f"| {self.image_size}px -> {self.vggt_input_size}px -> {self.target_grid}x{self.target_grid} " + f"tokens x {self.feature_dim}d" + ) + model.eval() + self.model = model + for param in self.model.parameters(): + param.requires_grad = False + + mean = torch.tensor(IMAGENET_DEFAULT_MEAN).view(1, 3, 1, 1) + std = torch.tensor(IMAGENET_DEFAULT_STD).view(1, 3, 1, 1) + self.register_buffer("_imagenet_mean", mean, persistent=False) + self.register_buffer("_imagenet_std", std, persistent=False) + + if self.norm_layer_type in ("bn", "ln"): + if self.norm_layer_type == "bn": + norm_builder = lambda: nn.SyncBatchNorm(self.feature_dim, affine=False) + else: + norm_builder = lambda: nn.LayerNorm(self.feature_dim, elementwise_affine=False) + self.latent_norms = nn.ModuleList([norm_builder() for _ in range(self.num_latent_layers)]) + else: + self.latent_norms = None + + def train(self, mode: bool = True): + # Keep the frozen backbone in eval mode. This is load-bearing beyond the + # usual dropout/BN reason: Aggregator gates `torch.utils.checkpoint` on + # `self.training`, which is pure waste under no_grad. + super().train(False) + self.model.eval() + return self + + def _denormalize_to_unit(self, images: torch.Tensor) -> torch.Tensor: + """Undo `video_aug`'s ImageNet normalization; the aggregator redoes it.""" + mean = self._imagenet_mean.to(device=images.device, dtype=images.dtype) + std = self._imagenet_std.to(device=images.device, dtype=images.dtype) + return (images * std + mean).clamp_(0.0, 1.0) + + def _to_vggt_resolution(self, images: torch.Tensor) -> torch.Tensor: + if images.shape[-2:] == (self.vggt_input_size, self.vggt_input_size): + return images + return F.interpolate( + images, + size=(self.vggt_input_size, self.vggt_input_size), + mode="bilinear", + align_corners=False, + ) + + def _pool_to_target_grid(self, tokens: torch.Tensor) -> torch.Tensor: + """[N, side*side, D] patch tokens -> [N, target_grid**2, D].""" + n, num_patches, dim = tokens.shape + side = int(math.sqrt(num_patches)) + if side * side != num_patches: + raise ValueError(f"VGGT patch count {num_patches} is not a square grid.") + if side == self.target_grid: + return tokens + grid = tokens.reshape(n, side, side, dim).permute(0, 3, 1, 2) + grid = F.adaptive_avg_pool2d(grid, (self.target_grid, self.target_grid)) + return grid.flatten(2).transpose(1, 2).contiguous() + + def _apply_norm(self, tokens: torch.Tensor, idx: int) -> torch.Tensor: + if not self.enable_norm: + return tokens + if self.norm_layer_type == "bn": + if self.latent_norms is None: + raise ValueError("VGGTEncoder has no BN layer initialized; set norm_layer_type to 'bn'.") + tokens_2d = tokens.reshape(-1, self.feature_dim) + tokens_2d = self.latent_norms[idx](tokens_2d) + return tokens_2d.view(tokens.shape[0], tokens.shape[1], self.feature_dim) + if self.norm_layer_type == "ln": + if self.latent_norms is None: + raise ValueError("VGGTEncoder has no LN layer initialized; set norm_layer_type to 'ln'.") + return self.latent_norms[idx](tokens) + if self.norm_layer_type == "l2": + return F.normalize(tokens, p=2, dim=-1) + return tokens + + @torch.no_grad() + def encode(self, images: torch.Tensor, remove_cls: bool = True, n: Union[int, Sequence] = -1) -> torch.Tensor: + if not remove_cls: + raise NotImplementedError( + "VGGTEncoder always strips camera/register tokens: the remaining patch tokens " + "are pooled onto a square grid, which the special tokens cannot join." + ) + if images.dim() == 5: + B, T = images.shape[0], images.shape[1] + flat = images.reshape(-1, images.shape[-3], images.shape[-2], images.shape[-1]) + elif images.dim() == 4: + B, T = images.shape[0], 1 + flat = images + else: + raise ValueError(f"Expected 4D or 5D input, got {tuple(images.shape)}") + + pixels = self._to_vggt_resolution(self._denormalize_to_unit(flat.float())) + param_dtype = next(self.model.parameters()).dtype + # S=1 -> every frame is encoded independently; see class docstring (3). + token_list, patch_start_idx = self.model.aggregator(pixels.unsqueeze(1).to(dtype=param_dtype)) + + cached = [t for t in token_list if t is not None] + if not cached: + raise RuntimeError("VGGT aggregator returned no cached layers.") + + list_n = [n] if isinstance(n, int) else list(n) + assert len(list_n) <= self.num_latent_layers, ( + f"VGGTEncoder expected at most {self.num_latent_layers} normalization layers, " + f"but received {len(list_n)} feature layers. Ensure this matches len(latent_layer_to_use)." + ) + + features = [] + for idx, layer in enumerate(list_n): + layer_tokens = cached[self._resolve_cached_index(layer, len(cached))] # [N, 1, 5+P, D] + layer_tokens = layer_tokens[:, 0, patch_start_idx:, :] # [N, P, D] + layer_tokens = self._pool_to_target_grid(layer_tokens) # [N, K, D] + layer_tokens = self._apply_norm(layer_tokens, idx) + features.append(layer_tokens.reshape(B, T, -1, self.feature_dim)) + + return features[0].detach() if isinstance(n, int) else [f.detach() for f in features] + + @staticmethod + def _resolve_cached_index(layer: int, num_cached: int) -> int: + """Map a config layer index onto the compacted list of cached layers. + + The configs speak DINOv3 (`-2` = penultimate block). VGGT only keeps four + blocks, so `-1` -> block 23 and `-2` -> block 17. Indices beyond the + cached range are clamped rather than silently returning `None`. + """ + layer = int(layer) + if layer < 0: + layer = max(layer, -num_cached) + else: + layer = min(layer, num_cached - 1) + return layer + + +def _load_vggt_state_dict(model_id: str) -> dict: + """Load VGGT weights from a local .pt/.safetensors file, a directory, or the hub.""" + path = _get_existing_path(model_id) + if path is not None: + if path.is_dir(): + candidates = [ + path / "model.pt", + path / "model.safetensors", + path / "pytorch_model.bin", + ] + found = next((c for c in candidates if c.exists()), None) + if found is None: + raise FileNotFoundError( + f"No VGGT weight file (model.pt / model.safetensors / pytorch_model.bin) under {path}" + ) + path = found + if path.suffix == ".safetensors": + from safetensors.torch import load_file + + return load_file(str(path)) + state_dict = _safe_torch_load(path) + return state_dict.get("model", state_dict) + + # Not a local path: fall back to the hub via PyTorchModelHubMixin. + from vggt.models.vggt import VGGT + + return VGGT.from_pretrained(model_id).state_dict() + + class CosmosAutoencoder(nn.Module): def __init__( @@ -610,7 +848,15 @@ def build_vision_encoder( key = str(model_id).lower() vjepa_hub_name = _infer_vjepa_hub_name(model_id) - if "dinov3" in key: + if "vggt" in key: + encoder = VGGTEncoder( + model_id=model_id, + num_latent_layers=num_latent_layers, + norm_layer_type=norm_layer_type, + enable_norm=enable_norm, + ) + return encoder, encoder.feature_dim + elif "dinov3" in key: encoder = DINOv3Encoder( model_id=model_id, num_latent_layers=num_latent_layers,