Model › Framework¶

@classmethod
def from_pretrained(
    cls,
    pretrained_checkpoint: str,
    **kwargs,
) -> None:
    """
    Restore a model instance from a saved checkpoint.

    Workflow:
        1. Resolve checkpoint path
        2. Load config + dataset normalization statistics
        3. Build model with loaded config
        4. Load state_dict strictly (reports missing/unexpected keys)
        5. Attach normalization stats for later un-normalization

    Args:
        pretrained_checkpoint: Path to .pt/.safetensors file or self-contained checkpoint directory.
        **kwargs: Extra constructor overrides passed to subclass.

    Returns:
        BaseFramework: Instantiated model (left on CPU; caller decides device).

    Raises:
        RuntimeError: If state_dict key mismatch occurs under strict=True.
        FileNotFoundError: If underlying files are missing (surfaced earlier).
    """
    pretrained_checkpoint = Path(pretrained_checkpoint)

    # lpt0309: 支持自包含目录格式checkpoint（单一路径推断，无需base_vlm）
    if pretrained_checkpoint.is_dir():
        logger.info(f"[lpt0309] Loading from self-contained checkpoint directory: {pretrained_checkpoint}")
        model_config, norm_stats = read_mode_config(pretrained_checkpoint)  # lpt0309: 从目录读取config和norm_stats

        config = dict_to_namespace(model_config)
        config.trainer.pretrained_checkpoint = None

        # 单次加载优化 - 如果checkpoint目录包含vlm_pretrained/（兼容旧名qwen_pretrained/），
        # 直接从中加载tokenizer/config，用meta device创建模型骨架，由后续load_state_dict一次性加载所有权重
        vlm_pretrained_dir = pretrained_checkpoint / "vlm_pretrained"
        legacy_dir = pretrained_checkpoint / "qwen_pretrained"
        if not (vlm_pretrained_dir.is_dir() and any(vlm_pretrained_dir.iterdir())):
            # 兼容旧格式 qwen_pretrained/
            if legacy_dir.is_dir() and any(legacy_dir.iterdir()):
                vlm_pretrained_dir = legacy_dir

        if vlm_pretrained_dir.is_dir() and any(vlm_pretrained_dir.iterdir()):
            logger.info(f"Found {vlm_pretrained_dir.name}/ in checkpoint, using single-read loading")
            cfg_key = _detect_vlm_cfg_key(config.framework)
            if cfg_key is not None:
                vlm_block = getattr(config.framework, cfg_key)
                original_base_vlm = getattr(vlm_block, 'base_vlm', "") if not hasattr(vlm_block, 'get') else vlm_block.get('base_vlm', "")
                vlm_block.vlm_type = original_base_vlm
                vlm_block.base_vlm = str(vlm_pretrained_dir)
                vlm_block._meta_device_init = True
            else:
                logger.warning("vlm_pretrained/ found but no VLM config key detected, skipping single-read optimization")
        else:
            logger.warning(f"No vlm_pretrained/ found (or empty), falling back to two-read loading from original base_vlm")

        FrameworkModel = build_framework(cfg=config)
        FrameworkModel.norm_stats = norm_stats

        # lpt0309: 从目录中找到权重文件
        weights_path = pretrained_checkpoint / "model.safetensors"
        if not weights_path.exists():
            weights_path = pretrained_checkpoint / "pytorch_model.pt"
        assert weights_path.exists(), f"[lpt0309] No weights file found in {pretrained_checkpoint}"

        if weights_path.suffix == ".safetensors":
            from safetensors.torch import load_file
            model_state_dict = load_file(str(weights_path))
        else:
            model_state_dict = torch.load(weights_path, map_location="cpu")

        logger.info(f"[lpt0309] Loading weights from {weights_path}")

        # Key remapping: old checkpoints use 'vlm.' prefix, new model uses 'qwen_vl_interface.'
        remapped = {}
        for k, v in model_state_dict.items():
            new_k = k.replace('vlm.', 'qwen_vl_interface.', 1) if k.startswith('vlm.') else k
            remapped[new_k] = v
        if len(remapped) != len(model_state_dict):
            logger.warning(f"Key remapping changed key count: {len(model_state_dict)} -> {len(remapped)}")
        else:
            n_remapped = sum(1 for k in model_state_dict if k.startswith('vlm.'))
            if n_remapped > 0:
                logger.info(f"Remapped {n_remapped} keys from vlm.* to qwen_vl_interface.*")
        model_state_dict = remapped

        model_keys = set(FrameworkModel.state_dict().keys())
        checkpoint_keys = set(model_state_dict.keys())
        # Try strict first; fall back to non-strict if only non-critical keys mismatch
        try:
            FrameworkModel.load_state_dict(model_state_dict, strict=True)
        except RuntimeError as e:
            common_keys = model_keys.intersection(checkpoint_keys)
            missing_keys = model_keys - common_keys
            unexpected_keys = checkpoint_keys - common_keys
            if missing_keys:
                logger.warning(f"Missing keys in state_dict ({len(missing_keys)}): {missing_keys}")
            if unexpected_keys:
                logger.warning(f"Unexpected keys in state_dict ({len(unexpected_keys)}): {unexpected_keys}")
            # Fall back to non-strict loading for cross-framework weight loading (e.g. openpi → AlphaBrain)
            logger.warning(f"Strict loading failed, falling back to non-strict (missing={len(missing_keys)}, unexpected={len(unexpected_keys)})")
            FrameworkModel.load_state_dict(model_state_dict, strict=False)

        logger.info(
            "[lpt0324] Successfully loaded model from self-contained checkpoint "
            "with legacy two-stage loading"
        )
        return FrameworkModel

    # origin0309: 原始文件格式加载（需要base_vlm路径，存在冗余权重读取）
    else:
        model_config, norm_stats = read_mode_config(pretrained_checkpoint)  # read config and norm_stats

        config = dict_to_namespace(model_config)
        model_config = config
        model_config.trainer.pretrained_checkpoint = None
        # FrameworkModel = cls(config=model_config, **kwargs) # TODO find cls by config
        FrameworkModel = build_framework(cfg=model_config)
        # set for action un-norm
        FrameworkModel.norm_stats = norm_stats
        # Load from Checkpoint (Custom --> should load both *projector* and *llm* weights)
        if pretrained_checkpoint.suffix == ".safetensors":
            from safetensors.torch import load_file
            # TODO pretrained_checkpoint 这里先转成了path后面又用str, 存在冗余
            model_state_dict = load_file(str(pretrained_checkpoint))
        else:
            model_state_dict = torch.load(pretrained_checkpoint, map_location="cpu")
        # logger.info(f"Loading model weights from `{pretrained_checkpoint}`")
        model_keys = set(FrameworkModel.state_dict().keys())
        checkpoint_keys = set(model_state_dict.keys())  # TODO 为什么会存在重复?
        try:
            FrameworkModel.load_state_dict(model_state_dict, strict=True)
        except RuntimeError as e:
            # must keep all keys matched
            common_keys = model_keys.intersection(checkpoint_keys)
            missing_keys = model_keys - common_keys
            unexpected_keys = checkpoint_keys - common_keys
            if missing_keys:
                logger.warning(f"Missing keys in state_dict: {missing_keys}")
            if unexpected_keys:
                logger.warning(f"Unexpected keys in state_dict: {unexpected_keys}")

            raise e

        # **ensure model is on GPU**
        FrameworkModel = FrameworkModel
        return FrameworkModel

@staticmethod
def convert_checkpoint_to_dir(
    old_ckpt_path: str,
    output_dir: str = None,
    base_vlm_path: str = None,
):
    """
    Convert an old-format file checkpoint to the new self-contained directory format.

    Args:
        old_ckpt_path: Path to old .safetensors/.pt checkpoint file.
        output_dir: Output directory path. If None, creates a directory alongside the file.
        base_vlm_path: Path to Qwen base model (for saving config + processor).
                      If None, reads from the checkpoint's config.yaml.
    """
    import shutil
    old_ckpt_path = Path(old_ckpt_path)
    assert old_ckpt_path.is_file(), f"Old checkpoint not found: {old_ckpt_path}"

    # Determine output directory
    if output_dir is None:
        output_dir = old_ckpt_path.parent / old_ckpt_path.stem.replace("_model", "").replace("_pytorch", "")
    output_dir = Path(output_dir)
    os.makedirs(output_dir, exist_ok=True)

    # Copy weights
    weights_name = "model.safetensors" if old_ckpt_path.suffix == ".safetensors" else "pytorch_model.pt"
    shutil.copy2(str(old_ckpt_path), str(output_dir / weights_name))
    logger.info(f"[lpt0309] Copied weights to {output_dir / weights_name}")

    # Copy config.yaml and dataset_statistics.json from run dir
    run_dir = old_ckpt_path.parents[1]
    for fname, target_name in [("config.yaml", "framework_config.yaml"), ("dataset_statistics.json", "dataset_statistics.json")]:
        src = run_dir / fname
        if src.exists():
            shutil.copy2(str(src), str(output_dir / target_name))
            logger.info(f"[lpt0309] Copied {fname} -> {output_dir / target_name}")

    # Save VLM config + processor (auto-detect VLM type from config)
    if base_vlm_path is None:
        config_yaml = run_dir / "config.yaml"
        if config_yaml.exists():
            from omegaconf import OmegaConf
            cfg = OmegaConf.load(str(config_yaml))
            base_vlm_path = _get_base_vlm_path(cfg.framework)

    if base_vlm_path:
        vlm_pretrained_dir = output_dir / "vlm_pretrained"
        os.makedirs(vlm_pretrained_dir, exist_ok=True)
        try:
            from transformers import AutoConfig, AutoProcessor
            vlm_config = AutoConfig.from_pretrained(base_vlm_path, trust_remote_code=True)
            vlm_config.save_pretrained(str(vlm_pretrained_dir))
            processor = AutoProcessor.from_pretrained(base_vlm_path)
            processor.save_pretrained(str(vlm_pretrained_dir))
            logger.info(f"[lpt0309] Saved VLM config + processor to {vlm_pretrained_dir}")
        except Exception as e:
            logger.warning(f"[lpt0309] Could not save VLM config/processor from {base_vlm_path}: {e}")

    logger.info(f"[lpt0309] Conversion complete: {output_dir}")

@classmethod
def get_action_stats(self, unnorm_key=None):
    """
    Retrieve raw action normalization statistics.

    Args:
        unnorm_key: Optional dataset stats key.

    Returns:
        dict: Stats structure (e.g. q01, q99, mask).
    """
    unnorm_key = self._check_unnorm_key(self.norm_stats, unnorm_key)
    return self.norm_stats[unnorm_key]["action"]

@staticmethod
def unnormalize_actions(normalized_actions: np.ndarray, action_norm_stats: Dict[str, np.ndarray]) -> np.ndarray:
    """
    Map normalized actions (≈[-1, 1]) back to original value range.

    Auto-detects normalization mode via the optional 'norm_mode' key in
    action_norm_stats (defaults to 'q99' for backward compatibility):
        - 'q99'     → uses q01 / q99 bounds
        - 'min_max' → uses min / max bounds

    Steps:
        - Clamp values to [-1, 1]
        - Threshold channel index 6 to {0,1} (binary semantic)
        - Apply linear scaling for masked dimensions

    Args:
        normalized_actions: Array shape [T, D] (or chunk length × action_dim).
        action_norm_stats: Dict containing stat arrays and optional 'norm_mode'.

    Returns:
        np.ndarray: Unnormalized actions (same shape as input).
    """
    norm_mode = action_norm_stats.get("norm_mode", "q99")
    if norm_mode == "min_max":
        ref_key_high, ref_key_low = "max", "min"
    else:
        ref_key_high, ref_key_low = "q99", "q01"
    mask = action_norm_stats.get("mask", np.ones_like(action_norm_stats[ref_key_low], dtype=bool))
    action_high = np.array(action_norm_stats[ref_key_high])
    action_low = np.array(action_norm_stats[ref_key_low])
    normalized_actions = np.clip(normalized_actions, -1, 1)
    normalized_actions[:, 6] = np.where(normalized_actions[:, 6] < 0.5, 0, 1)
    actions = np.where(
        mask,
        0.5 * (normalized_actions + 1) * (action_high - action_low) + action_low,
        normalized_actions,
    )

    return actions

def get(self, key, default=None):
    """
    Return attribute value if present, else default (dict-like API).

    Args:
        key: Attribute name.
        default: Fallback if attribute missing.

    Returns:
        Any: Stored value or default.
    """
    return getattr(self, key, default)

def items(self):
    """
    Iterate (key, value) pairs like dict.items().

    Returns:
        Generator[Tuple[str, Any], None, None]
    """
    return ((key, getattr(self, key)) for key in self.__dict__)

def to_dict(self):
    """
    Recursively convert nested NamespaceWithGet objects into plain dicts.

    Returns:
        dict: Fully materialized dictionary structure.
    """
    return {key: value.to_dict() if isinstance(value, NamespaceWithGet) else value for key, value in self.items()}

def forward(self, x: torch.Tensor) -> torch.Tensor:
    """x: (B, 3, H, W) → (B, img_feat_dim)"""
    x = self.net(x)
    x = x.flatten(1)
    return self.proj(x)

def forward(self, texts: List[str]) -> torch.Tensor:
    """texts: List[str] → (B, text_feat_dim)"""
    # 把文本 hash 成 vocab_size 以内的 index，简单但可分辨
    indices = []
    for t in texts:
        h = int(hashlib.md5(t.encode()).hexdigest(), 16) % self.vocab_size
        indices.append(h)
    idx = torch.tensor(indices, device=self.emb.weight.device)
    return self.emb(idx)  # (B, text_feat_dim)

def forward(self, examples: List[dict], **kwargs) -> dict:
    """
    Returns:
        {"action_loss": scalar tensor}
    """
    fused = self._encode_batch(examples)                      # (B, fuse_dim)
    raw   = self.action_head(fused)                           # (B, chunk_len*action_dim)
    pred  = raw.view(-1, self.chunk_len, self.action_dim)     # (B, chunk_len, action_dim)

    # 标签
    actions = np.array([ex["action"] for ex in examples])    # (B, T, D)
    tgt = torch.tensor(actions, device=pred.device, dtype=pred.dtype)
    # 取最后 chunk_len 步
    tgt = tgt[:, -self.chunk_len:, :]                         # (B, chunk_len, action_dim)

    loss = self.l1_loss(pred, tgt)
    return {"action_loss": loss}

def predict_action(self, examples: List[dict], **kwargs) -> dict:
    """
    Returns:
        {"normalized_actions": np.ndarray (B, chunk_len, action_dim)}
    """
    with torch.no_grad():
        fused = self._encode_batch(examples)
        raw   = self.action_head(fused)
        pred  = raw.view(-1, self.chunk_len, self.action_dim)
    return {"normalized_actions": pred.cpu().float().numpy()}

def forward(self, state: torch.Tensor, action_chunk: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
    """
    state:        (B, state_dim)
    action_chunk: (B, chunk_len, action_dim)
    Returns: mu, log_var  each (B, latent_dim)
    """
    B = state.shape[0]
    state_emb  = self.state_proj(state).unsqueeze(1)       # (B, 1, H)
    action_emb = self.action_proj(action_chunk)             # (B, T, H)
    cls = self.cls_token.expand(B, -1, -1)                 # (B, 1, H)
    seq = torch.cat([cls, state_emb, action_emb], dim=1)   # (B, 2+T, H)
    out = self.encoder(seq)                                 # (B, 2+T, H)
    cls_out = out[:, 0, :]                                  # (B, H)
    return self.mu_head(cls_out), self.var_head(cls_out)