data.py

#!/usr/bin/env python3
"""
data.py  (stage3 / Flow Matching — Pi0.5 风格)
==============================================
多数据集 LeRobot 格式数据加载器 (兼容 v2.0 和 v3.0)，支持：
  - 多数据集加权混合采样 (BridgeData / DROID / Libero / 自有数据)
  - 机器人状态 (proprioception) 读取与归一化
  - Action chunking (K=16)
  - 多相机输入 (base_rgb + wrist_rgb)
  - 域随机化 (Domain Randomization) 缓解背景过拟合
  - 标准 LeRobot 命名 与 HuggingFace LeRobot 下载的 "data__chunk__file" 命名兼容
  - 缺视频时自动回退到占位图，避免训练中断

v2.0 vs v3.0:
  v2.0: 1 parquet = 1 episode, 视频 episode_*.mp4, task_index + tasks.jsonl
  v3.0: 1 parquet = N episodes (is_first/is_last 分界), 视频 file-*.mp4,
        多 camera keys, language_instruction 列直接读指令

接口契约
========
build_dataloader(config) -> (train_loader, val_loader)

每个 batch 必须包含:
{
    "image":        List[PIL.Image] 或 List[List[PIL.Image]],
    "instruction":  List[str],
    "state":        torch.Tensor,         # (B, state_dim)  或 None
    "action":       torch.Tensor,         # (B, action_dim)
    "action_chunk": torch.Tensor,         # (B, K, action_dim)
    "raw_action":   torch.Tensor,         # (B, action_dim)
    "frame_index":  List[int],
    "episode_index": List[int],
    "dataset_name": List[str],            # 可选
}

config 示例见 config_flow.yaml；data.py 新增字段:
  data:
    datasets:
      - name: bridge
        dir: /mnt/workspace/Dataset/stage_frozen01/full
        weight: 0.4
      - name: droid
        dir: /mnt/workspace/Dataset/droid/full
        weight: 0.2
        camera_keys: [observation.images.exterior_1_left, ...]
    val_ratio: 0.05
    action_chunk_size: 16
    action_dim: 7
    state_dim: 7
    domain_randomization: true
"""

from __future__ import annotations

import json
import glob
import os
import random
import threading
from pathlib import Path
from typing import Any, Callable, Dict, List, Optional, Tuple, Union

import numpy as np
import pandas as pd
import torch
from PIL import Image
from torch.utils.data import (
    ConcatDataset, DataLoader, Dataset, WeightedRandomSampler,
)
try:
    from .utils.chunking import TemporalEnsembler
    HAS_CHUNKING = True
except ImportError:
    HAS_CHUNKING = False

try:
    from torchvision import transforms
except ModuleNotFoundError:
    transforms = None  # type: ignore

# 视频解码 (优先 decord，否则 imageio)
try:
    import decord
    decord.bridge.set_bridge("torch")
    HAS_DECORD = True
except Exception:
    HAS_DECORD = False
    try:
        import imageio
    except Exception:
        imageio = None

logger = __import__("logging").getLogger("LingArm.stage3")


# ---------------------------------------------------------------------------
# 视频帧缓存
# ---------------------------------------------------------------------------
class VideoCache:
    """按 episode 缓存解码后的视频帧 (LRU)。"""

    def __init__(self, max_cached_episodes: int = 200):
        self.cache: Dict[str, np.ndarray] = {}
        self.access_order: List[str] = []
        self.max_cached = max_cached_episodes
        self.lock = threading.Lock()

    def __getstate__(self):
        state = self.__dict__.copy()
        state.pop("lock", None)
        state["cache"] = {}  # 不跨进程传输大缓存
        state["access_order"] = []
        return state

    def __setstate__(self, state):
        self.__dict__.update(state)
        self.lock = threading.Lock()

    def get_frame(
        self,
        video_path: str,
        frame_idx: int,
    ) -> np.ndarray:
        with self.lock:
            if video_path not in self.cache:
                self._load_video(video_path)
            else:
                self.access_order.remove(video_path)
                self.access_order.append(video_path)
            frames = self.cache[video_path]
            if frame_idx >= len(frames):
                frame_idx = len(frames) - 1
            return frames[frame_idx]

    def _load_video(self, video_path: str):
        while len(self.cache) >= self.max_cached and self.access_order:
            old = self.access_order.pop(0)
            self.cache.pop(old, None)

        if not os.path.exists(video_path):
            raise FileNotFoundError(f"Video not found: {video_path}")

        if HAS_DECORD:
            vr = decord.VideoReader(video_path)
            frames = vr.get_batch(list(range(len(vr))))
            # decord 0.6.0+ 返回 torch.Tensor, 旧版返回 NDArray with .asnumpy()
            if hasattr(frames, 'numpy'):
                self.cache[video_path] = frames.numpy()
            else:
                self.cache[video_path] = frames.asnumpy()
        elif imageio is not None:
            reader = imageio.get_reader(video_path)
            self.cache[video_path] = np.array([
                reader.get_data(i) for i in range(reader.count_frames())
            ])
            reader.close()
        else:
            raise RuntimeError(
                "Need decord or imageio to read videos. "
                "Install: pip install decord"
            )
        self.access_order.append(video_path)


# ---------------------------------------------------------------------------
# 缺视频时的占位图
# ---------------------------------------------------------------------------
class DummyImageCache:
    """为没有视频的样本生成一张随机纹理占位图。"""

    def __init__(self, size: int = 224):
        self.size = size
        self._cache: Dict[Tuple[int, int], Image.Image] = {}

    def get(self, seed: int = 0) -> Image.Image:
        key = (self.size, seed)
        if key not in self._cache:
            rng = np.random.default_rng(seed)
            arr = rng.integers(0, 255, (self.size, self.size, 3), dtype=np.uint8)
            self._cache[key] = Image.fromarray(arr).convert("RGB")
        return self._cache[key]


DUMMY_IMAGE = DummyImageCache()


# ---------------------------------------------------------------------------
# LeRobot 数据集
# ---------------------------------------------------------------------------
class LeRobotDataset(Dataset):
    """
    LeRobot 格式数据集加载器 (stage3), 兼容 v2.0 和 v3.0。

    v2.0: 1 parquet = 1 episode, 视频 episode_*.mp4, task_index → tasks.jsonl
    v3.0: 1 parquet = N episodes (is_first/is_last 分界),
          视频 file-*.mp4, videos/{cam}/chunk-*/ 路径,
          language_instruction 列直接读取任务指令

    Args:
        data_dir: 数据集根目录 (含 meta/info.json)
        dataset_name: 数据集名称 (用于日志/加权)
        action_dim: 目标动作维度 (默认 7)
        state_dim: 目标状态维度 (默认 7)
        action_chunk_size: K (默认 16)
        image_size: 图像目标尺寸 (PIL resize；最终由 processor 决定)
        action_normalize: 是否 z-score 标准化 action
        state_normalize: 是否 z-score 标准化 state 后 clip 到 [-1,1]
        state_low/state_high: state clip 范围
        frame_sampling: "all" | "uniform"
        max_frames_per_episode: uniform 采样时最大帧数
        use_processor: True 返回 PIL.Image，False 返回 tensor
        camera_keys: 多相机 key 列表
        dummy_image_on_missing_video: 缺视频是否用占位图
        augment: 是否启用域随机化
        action_adapter: 可选动作转换函数
        state_adapter: 可选状态转换函数
        val_episode_ids: 若提供，仅使用这些 episode (验证集)
    """

    def __init__(
        self,
        data_dir: str,
        dataset_name: str = "lerobot",
        action_dim: int = 7,
        state_dim: int = 7,
        action_chunk_size: int = 16,
        image_size: int = 224,
        action_normalize: bool = True,
        state_normalize: bool = True,
        state_low: float = -1.0,
        state_high: float = 1.0,
        frame_sampling: str = "all",
        max_frames_per_episode: Optional[int] = None,
        use_processor: bool = True,
        max_cache_episodes: int = 200,
        camera_keys: Optional[List[str]] = None,
        dummy_image_on_missing_video: bool = True,
        augment: bool = True,
        action_adapter: Optional[Callable[[np.ndarray, str], np.ndarray]] = None,
        state_adapter: Optional[Callable[[np.ndarray, str], np.ndarray]] = None,
        val_episode_ids: Optional[set] = None,
        skip_stats_compute: Optional[bool] = None,
    ):
        self.data_dir = Path(data_dir)
        self.dataset_name = dataset_name
        self.action_dim = action_dim
        self.state_dim = state_dim
        self.action_chunk_size = action_chunk_size
        self.image_size = image_size
        self.action_normalize = action_normalize
        self.state_normalize = state_normalize
        self.state_low = state_low
        self.state_high = state_high
        self.frame_sampling = frame_sampling
        self.max_frames_per_episode = max_frames_per_episode
        self.use_processor = use_processor
        self.camera_keys = camera_keys or ["observation.images.image_0"]
        self.dummy_image_on_missing_video = dummy_image_on_missing_video
        self.augment = augment
        self.action_adapter = action_adapter
        self.state_adapter = state_adapter
        self.val_episode_ids = val_episode_ids

        # 检测 LeRobot 格式版本 (v2.0 / v3.0)
        self.codebase_version = self._detect_version()
        logger.info("[Dataset %s] codebase_version=%s", dataset_name, self.codebase_version)

        # 读取 meta / 任务指令
        self.tasks = self._load_tasks()

        # 构建 episode 索引
        self.episodes = self._build_index()
        if not self.episodes:
            raise RuntimeError(f"No valid episodes found in {data_dir}")

        # 计算统计量 (仅训练集)
        if skip_stats_compute is None:
            skip_stats_compute = (val_episode_ids is not None)
        if skip_stats_compute:
            self.action_mean = self.action_std = None
            self.state_mean = self.state_std = None
        else:
            self.action_mean, self.action_std = self._compute_action_stats()
            self.state_mean, self.state_std = self._compute_state_stats()

        # 视频缓存
        self.video_cache = VideoCache(max_cached_episodes=max_cache_episodes)

        # 图像预处理
        if not use_processor:
            if transforms is None:
                raise ImportError("torchvision is required when use_processor=False")
            self.image_transform = transforms.Compose([
                transforms.Resize((image_size, image_size)),
                transforms.ToTensor(),
                transforms.Normalize(
                    mean=[0.485, 0.456, 0.406],
                    std=[0.229, 0.224, 0.225],
                ),
            ])
        else:
            self.image_transform = None

        # 域随机化 (仅对 PIL 做轻量变换)
        self.aug_transform = None
        if augment and transforms is not None:
            self.aug_transform = transforms.Compose([
                transforms.ColorJitter(brightness=0.15, contrast=0.15, saturation=0.15, hue=0.05),
                transforms.RandomGrayscale(p=0.05),
            ])

        # 镜像翻转增广 (保持关闭, 镜像几何变换需要逐相机标定)
        self.hflip_aug = False
        self.hflip_prob = 0.5

        # 随机透视抖动 (模拟相机位置变化)
        self.perspective_aug = augment
        self.perspective_distortion = 0.15

        total_frames = sum(len(ep["frame_indices"]) for ep in self.episodes)
        logger.info(
            "[Dataset %s] episodes=%d frames=%d action_dim=%d state_dim=%d cameras=%d",
            dataset_name, len(self.episodes), total_frames,
            self.action_dim, self.state_dim, len(self.camera_keys),
        )

    # ------------------------------------------------------------------
    # 版本检测
    # ------------------------------------------------------------------
    def _detect_version(self) -> str:
        """从 meta/info.json 读取 codebase_version，回退到 v2.0。"""
        for candidate in ["meta/info.json", "info.json"]:
            info_path = self.data_dir / candidate
            if info_path.exists():
                try:
                    with open(info_path, "r", encoding="utf-8") as f:
                        info = json.load(f)
                    return info.get("codebase_version", "v2.0")
                except Exception:
                    pass
        return "v2.0"

    # ------------------------------------------------------------------
    # Meta / 任务指令
    # ------------------------------------------------------------------
    def _load_tasks(self) -> Dict[int, str]:
        """加载任务指令表。v3.0 返回空 dict (指令从 parquet language_instruction 列读取)。"""
        if self.codebase_version == "v3.0":
            return {}

        tasks: Dict[int, str] = {}

        # 1) meta/tasks.jsonl
        jsonl_path = self.data_dir / "meta" / "tasks.jsonl"
        if jsonl_path.exists():
            with open(jsonl_path, "r", encoding="utf-8") as f:
                for line in f:
                    line = line.strip()
                    if not line:
                        continue
                    obj = json.loads(line)
                    tasks[obj["task_index"]] = obj.get("task", "")
            return tasks

        # 2) meta__tasks.parquet (HF LeRobot 下载格式)
        parquet_paths = list(self.data_dir.glob("meta__tasks*.parquet"))
        if not parquet_paths:
            parquet_paths = list(self.data_dir.glob("*/meta__tasks*.parquet"))
        for p in parquet_paths:
            try:
                df = pd.read_parquet(p)
                if "task_index" not in df.columns:
                    continue
                has_task_col = "task" in df.columns
                for idx, row in df.iterrows():
                    ti = int(row["task_index"])
                    task_text = str(row["task"]) if has_task_col else (
                        str(idx) if isinstance(idx, str) else "")
                    tasks[ti] = task_text
                return tasks
            except Exception as e:
                logger.warning("[Dataset %s] failed to read task parquet %s: %s",
                               self.dataset_name, p, e)

        return tasks

    # ------------------------------------------------------------------
    # 索引构建
    # ------------------------------------------------------------------
    def _find_parquet_files(self) -> List[Path]:
        patterns = [
            "data/chunk-*/episode_*.parquet",          # 标准 v2.0
            "data/chunk-*/file-*.parquet",             # 标准 v3.0 (DROID)
            "data__chunk-*__file-*.parquet",           # HF 下载格式
            "*/data__chunk-*__file-*.parquet",         # 嵌套 HF 下载
        ]
        files: List[Path] = []
        for pat in patterns:
            files.extend(self.data_dir.glob(pat))
        return sorted(set(files))

    def _build_index(self) -> List[Dict[str, Any]]:
        episodes = []
        for pf in self._find_parquet_files():
            try:
                df = pd.read_parquet(pf)
            except Exception as e:
                logger.warning("[Dataset %s] skip parquet %s: %s", self.dataset_name, pf, e)
                continue

            if "action" not in df.columns:
                logger.warning("[Dataset %s] skip %s: no action column", self.dataset_name, pf)
                continue

            if len(df) == 0:
                continue

            if self.codebase_version == "v3.0":
                episodes.extend(self._build_index_v3(pf, df))
            else:
                episodes.extend(self._build_index_v2(pf, df))

        return episodes

    def _build_index_v2(self, pf: Path, df: pd.DataFrame) -> List[Dict[str, Any]]:
        """v2.0: 1 parquet = 1 episode。"""
        n_frames = len(df)

        if "episode_index" in df.columns:
            ep_idx = int(df["episode_index"].iloc[0])
        else:
            ep_idx = 0

        if self.val_episode_ids is not None and ep_idx not in self.val_episode_ids:
            return []

        chunk_name, video_stem = self._video_location(pf)
        video_paths_per_cam, has_any_video = self._resolve_video_paths_v2(chunk_name, video_stem)

        task_index = int(df["task_index"].iloc[0]) if "task_index" in df.columns else 0
        instruction = self.tasks.get(task_index, "")

        frame_indices = self._sample_frame_indices(n_frames, range(n_frames))
        raw_frame_indices = (
            df["frame_index"].tolist() if "frame_index" in df.columns
            else list(range(n_frames))
        )

        actions_raw = self._extract_actions(df)
        states_raw = self._extract_states(df)

        return [{
            "parquet": pf,
            "video_paths": video_paths_per_cam,
            "has_video": has_any_video,
            "n_frames": n_frames,
            "frame_indices": frame_indices,
            "instruction": instruction,
            "task_index": task_index,
            "episode_index": ep_idx,
            "actions_raw": actions_raw,
            "states_raw": states_raw,
            "raw_frame_indices": raw_frame_indices,
            "frame_offset": 0,  # v2.0: 无偏移 (1 episode = 1 mp4)
        }]

    def _build_index_v3(self, pf: Path, df: pd.DataFrame) -> List[Dict[str, Any]]:
        """v3.0: 1 parquet = N episodes, 按 is_first/is_last 拆分。"""
        if "is_first" not in df.columns or "is_last" not in df.columns:
            logger.warning(
                "[Dataset %s] v3.0 parquet lacks is_first/is_last: %s, "
                "treating as single episode", self.dataset_name, pf
            )
            return self._build_index_v2(pf, df)

        n_total = len(df)
        first_rows = df.index[df["is_first"] == True].tolist()
        last_rows = df.index[df["is_last"] == True].tolist()

        if not first_rows:
            first_rows = [0]
        if not last_rows:
            last_rows = [n_total - 1]
        if len(last_rows) < len(first_rows):
            last_rows.append(n_total - 1)

        chunk_name, video_stem = self._video_location(pf)
        video_paths_per_cam, has_any_video = self._resolve_video_paths_v3(chunk_name, video_stem)

        episodes = []
        for ep_i, (start_row, end_row) in enumerate(zip(first_rows, last_rows)):
            if end_row < start_row:
                continue
            ep_df = df.loc[start_row:end_row]
            n_frames = len(ep_df)
            if n_frames == 0:
                continue

            # episode_index
            if "episode_index" in ep_df.columns:
                ep_idx = int(ep_df["episode_index"].iloc[0])
            else:
                ep_idx = ep_i

            if self.val_episode_ids is not None and ep_idx not in self.val_episode_ids:
                continue

            # v3.0: 语言指令从 language_instruction 列读取
            if "language_instruction" in ep_df.columns:
                instruction = str(ep_df["language_instruction"].iloc[0])
            else:
                instruction = ""

            # 帧采样 (0-based within episode)
            frame_indices = self._sample_frame_indices(n_frames, range(n_frames))

            # 原始 frame_index (parquet 中的绝对帧号 → mp4 内的帧位置)
            if "frame_index" in ep_df.columns:
                raw_frame_indices = ep_df["frame_index"].tolist()
            else:
                raw_frame_indices = list(range(n_frames))

            frame_offset = int(raw_frame_indices[0]) if raw_frame_indices else 0

            actions_raw = self._extract_actions(ep_df)
            states_raw = self._extract_states(ep_df)

            episodes.append({
                "parquet": pf,
                "video_paths": video_paths_per_cam,
                "has_video": has_any_video,
                "n_frames": n_frames,
                "frame_indices": frame_indices,
                "instruction": instruction,
                "task_index": 0,
                "episode_index": ep_idx,
                "actions_raw": actions_raw,
                "states_raw": states_raw,
                "raw_frame_indices": raw_frame_indices,
                "frame_offset": frame_offset,  # mp4 内起始帧号
            })
        return episodes

    # ------------------------------------------------------------------
    # 视频路径解析
    # ------------------------------------------------------------------
    def _video_location(self, parquet_path: Path) -> Tuple[str, str]:
        """从 parquet 路径推断 chunk 名和视频 stem。"""
        stem = parquet_path.stem  # e.g. episode_000000 or file-000 or data__chunk-000__file-000
        if stem.startswith("data__"):
            parts = stem.split("__")
            return parts[1], parts[2]  # chunk-000, file-000
        else:
            chunk_name = parquet_path.parent.name  # chunk-xxx
            return chunk_name, stem

    def _resolve_video_paths_v2(
        self, chunk_name: str, video_stem: str
    ) -> Tuple[Dict[str, Optional[str]], bool]:
        """v2.0 视频路径: videos/<chunk>/<cam>/<stem>.mp4"""
        video_paths: Dict[str, Optional[str]] = {}
        has_any = False
        for cam_key in self.camera_keys:
            cam_dir = cam_key.replace(".", "_") if "." in cam_key else cam_key
            candidates = [
                self.data_dir / "videos" / chunk_name / cam_key / f"{video_stem}.mp4",
                self.data_dir / "videos" / chunk_name / cam_dir / f"{video_stem}.mp4",
            ]
            found = None
            for vp in candidates:
                if vp.exists():
                    found = str(vp)
                    break
            video_paths[cam_key] = found
            if found:
                has_any = True
        return video_paths, has_any

    def _resolve_video_paths_v3(
        self, chunk_name: str, video_stem: str
    ) -> Tuple[Dict[str, Optional[str]], bool]:
        """v3.0 视频路径: videos/<cam_key>/<chunk>/<stem>.mp4

        DROID v3.0 使用 videos/{video_key}/chunk-*/file-*.mp4 格式。
        尝试多种路径以兼容不同下载布局。
        """
        video_paths: Dict[str, Optional[str]] = {}
        has_any = False
        for cam_key in self.camera_keys:
            candidates = [
                # (1) v3 标准: videos/<cam>/<chunk>/<file>.mp4
                self.data_dir / "videos" / cam_key / chunk_name / f"{video_stem}.mp4",
                # (2) v2 风格回退: videos/<chunk>/<cam>/<file>.mp4
                self.data_dir / "videos" / chunk_name / cam_key / f"{video_stem}.mp4",
                # (3) 下划线版本
                self.data_dir / "videos" / chunk_name / cam_key.replace(".", "_") / f"{video_stem}.mp4",
            ]
            found = None
            for vp in candidates:
                if vp.exists():
                    found = str(vp)
                    break
            video_paths[cam_key] = found
            if found:
                has_any = True
        return video_paths, has_any

    # ------------------------------------------------------------------
    # 辅助: 帧采样 / action 提取 / state 提取
    # ------------------------------------------------------------------
    def _sample_frame_indices(
        self, n_frames: int, src_indices: Any
    ) -> List[int]:
        src = list(src_indices)
        if self.frame_sampling == "uniform" and self.max_frames_per_episode:
            if n_frames <= self.max_frames_per_episode:
                return src
            else:
                return np.linspace(0, n_frames - 1, self.max_frames_per_episode, dtype=int).tolist()
        return src

    def _extract_actions(self, df: pd.DataFrame) -> np.ndarray:
        arr = np.array(df["action"].tolist(), dtype=np.float32)
        if self.action_adapter is not None:
            arr = np.stack([self.action_adapter(a, self.dataset_name) for a in arr])
        else:
            arr = self._adapt_dim_batch(arr, self.action_dim, "action")
        return arr

    def _extract_states(self, df: pd.DataFrame) -> Optional[np.ndarray]:
        if "observation.state" not in df.columns:
            return None
        arr = np.array(df["observation.state"].tolist(), dtype=np.float32)
        if self.state_adapter is not None:
            arr = np.stack([self.state_adapter(s, self.dataset_name) for s in arr])
        else:
            arr = self._adapt_dim_batch(arr, self.state_dim, "state")
        return arr

    # ------------------------------------------------------------------
    # 统计量
    # ------------------------------------------------------------------
    def _compute_action_stats(self) -> Tuple[np.ndarray, np.ndarray]:
        """批量 Welford 在线算法。"""
        dim = self.action_dim
        count = 0
        mean = np.zeros(dim, dtype=np.float64)
        m2 = np.zeros(dim, dtype=np.float64)

        for ep in self.episodes:
            acts = ep["actions_raw"].astype(np.float64)
            n = len(acts)
            if n == 0:
                continue
            if count == 0:
                mean = acts.mean(axis=0)
                m2 = acts.var(axis=0) * n
                count = n
            else:
                batch_mean = acts.mean(axis=0)
                batch_var = acts.var(axis=0)
                delta = batch_mean - mean
                new_count = count + n
                new_mean = mean + delta * n / new_count
                m2 = m2 + batch_var * n + (delta ** 2) * count * n / new_count
                mean = new_mean
                count = new_count

        if count < 2:
            return np.zeros(dim, dtype=np.float32), np.ones(dim, dtype=np.float32)

        mean = mean.astype(np.float32)
        std = np.sqrt(m2 / count).astype(np.float32)
        std = np.where(std < 1e-6, 1.0, std)
        return mean, std

    def _compute_state_stats(self) -> Tuple[Optional[np.ndarray], Optional[np.ndarray]]:
        """批量 Welford 在线算法计算 state 统计。"""
        dim = self.state_dim
        count = 0
        mean = np.zeros(dim, dtype=np.float64)
        m2 = np.zeros(dim, dtype=np.float64)

        for ep in self.episodes:
            if ep["states_raw"] is None:
                continue
            states = ep["states_raw"].astype(np.float64)
            n = len(states)
            if n == 0:
                continue
            if count == 0:
                mean = states.mean(axis=0)
                m2 = states.var(axis=0) * n
                count = n
            else:
                batch_mean = states.mean(axis=0)
                batch_var = states.var(axis=0)
                delta = batch_mean - mean
                new_count = count + n
                new_mean = mean + delta * n / new_count
                m2 = m2 + batch_var * n + (delta ** 2) * count * n / new_count
                mean = new_mean
                count = new_count

        if count < 2:
            return None, None

        mean = mean.astype(np.float32)
        std = np.sqrt(m2 / count).astype(np.float32)
        std = np.where(std < 1e-6, 1.0, std)
        return mean, std

    def set_stats(
        self,
        action_mean: Optional[np.ndarray],
        action_std: Optional[np.ndarray],
        state_mean: Optional[np.ndarray],
        state_std: Optional[np.ndarray],
    ):
        """验证集复用训练集统计量。"""
        self.action_mean = action_mean
        self.action_std = action_std
        self.state_mean = state_mean
        self.state_std = state_std

    # ------------------------------------------------------------------
    # 适配到目标维度
    # ------------------------------------------------------------------
    def _adapt_action(self, raw: np.ndarray) -> np.ndarray:
        if self.action_adapter is not None:
            return self.action_adapter(raw, self.dataset_name)
        return self._pad_or_truncate(raw, self.action_dim, "action")

    def _adapt_state(self, raw: np.ndarray) -> np.ndarray:
        if self.state_adapter is not None:
            return self.state_adapter(raw, self.dataset_name)
        return self._pad_or_truncate(raw, self.state_dim, "state")

    def _pad_or_truncate(self, vec: np.ndarray, target_dim: int, name: str) -> np.ndarray:
        vec = np.asarray(vec, dtype=np.float32)
        d = vec.shape[0]
        if d == target_dim:
            return vec
        if d > target_dim:
            logger.warning(
                "[Dataset %s] %s dim %d > target %d, truncating first %d dims",
                self.dataset_name, name, d, target_dim, target_dim,
            )
            return vec[:target_dim]
        logger.warning(
            "[Dataset %s] %s dim %d < target %d, zero-padding",
            self.dataset_name, name, d, target_dim,
        )
        out = np.zeros(target_dim, dtype=np.float32)
        out[:d] = vec
        return out

    def _adapt_dim_batch(self, arr: np.ndarray, target_dim: int, name: str) -> np.ndarray:
        """批量 pad/truncate 到目标维度。"""
        d = arr.shape[1]
        if d == target_dim:
            return arr
        if d > target_dim:
            logger.warning(
                "[Dataset %s] %s dim %d > target %d, truncating (batch)",
                self.dataset_name, name, d, target_dim,
            )
            return arr[:, :target_dim]
        logger.warning(
            "[Dataset %s] %s dim %d < target %d, zero-padding (batch)",
            self.dataset_name, name, d, target_dim,
        )
        pad = np.zeros((arr.shape[0], target_dim - d), dtype=arr.dtype)
        return np.concatenate([arr, pad], axis=1)

    # ------------------------------------------------------------------
    # 图像读取
    # ------------------------------------------------------------------
    def _read_images(self, ep: Dict[str, Any], frame_idx: int) -> List[Image.Image]:
        images: List[Image.Image] = []
        seed = int(ep["episode_index"] * 100000 + frame_idx)
        # v3.0: 多 episode 共用 mp4, 实际帧号 = episode 内偏移 + mp4 起始偏移
        actual_frame = frame_idx + ep.get("frame_offset", 0)
        for cam_key in self.camera_keys:
            vp = ep["video_paths"].get(cam_key)
            if vp and os.path.exists(vp):
                frame = self.video_cache.get_frame(vp, actual_frame)
                img = Image.fromarray(frame).convert("RGB")
            elif self.dummy_image_on_missing_video:
                img = DUMMY_IMAGE.get(seed).copy()
                seed += 1
            else:
                raise FileNotFoundError(f"Missing video for camera {cam_key}: {vp}")

            if img.size[0] != self.image_size or img.size[1] != self.image_size:
                img = img.resize((self.image_size, self.image_size), Image.BILINEAR)

            if self.augment and self.aug_transform is not None and random.random() < 0.5:
                img = self.aug_transform(img)

            if self.perspective_aug and random.random() < 0.3:
                import torchvision.transforms.functional as TF
                w, h = img.size
                max_off = int(w * self.perspective_distortion)
                startpoints = [[0, 0], [w, 0], [w, h], [0, h]]
                endpoints = [
                    [random.randint(-max_off, max_off), random.randint(-max_off, max_off)],
                    [w + random.randint(-max_off, max_off), random.randint(-max_off, max_off)],
                    [w + random.randint(-max_off, max_off), h + random.randint(-max_off, max_off)],
                    [random.randint(-max_off, max_off), h + random.randint(-max_off, max_off)],
                ]
                img = TF.perspective(img, startpoints, endpoints, interpolation=Image.BILINEAR)

            images.append(img)

        return images

    # ------------------------------------------------------------------
    # Dataset 接口
    # ------------------------------------------------------------------
    def __len__(self) -> int:
        return sum(len(ep["frame_indices"]) for ep in self.episodes)

    def __getitem__(self, idx: int) -> Dict[str, Any]:
        cumulative = 0
        for ep in self.episodes:
            n = len(ep["frame_indices"])
            if idx < cumulative + n:
                frame_idx = ep["frame_indices"][idx - cumulative]
                break
            cumulative += n
        else:
            raise IndexError(f"Index {idx} out of range")

        images = self._read_images(ep, frame_idx)
        instruction = ep["instruction"]

        do_hflip = False
        if self.hflip_aug and random.random() < self.hflip_prob:
            do_hflip = True
            images = [img.transpose(Image.FLIP_LEFT_RIGHT) for img in images]
            instruction = instruction.replace("left", "\x00L\x00").replace("right", "left").replace("\x00L\x00", "right")

        raw_action = ep["actions_raw"][frame_idx].copy()
        if do_hflip and raw_action.shape[0] >= 6:
            raw_action[0] = -raw_action[0]
            raw_action[1] = -raw_action[1]
            raw_action[3] = -raw_action[3]
            raw_action[4] = -raw_action[4]

        norm_action = raw_action.copy()
        if self.action_normalize and self.action_mean is not None:
            norm_action = (raw_action - self.action_mean) / self.action_std

        if ep["states_raw"] is not None:
            raw_state = ep["states_raw"][frame_idx].copy()
            if do_hflip and raw_state.shape[0] >= 6:
                raw_state[0] = -raw_state[0]
                raw_state[1] = -raw_state[1]
                if raw_state.shape[0] >= 6:
                    raw_state[3] = -raw_state[3]
                    raw_state[4] = -raw_state[4]
            norm_state = raw_state.copy()
            if self.state_normalize and self.state_mean is not None:
                norm_state = (raw_state - self.state_mean) / self.state_std
                norm_state = np.clip(norm_state, self.state_low, self.state_high)
        else:
            norm_state = np.zeros(self.state_dim, dtype=np.float32)

        ep_actions = ep["actions_raw"]
        if self.action_normalize and self.action_mean is not None:
            ep_actions = (ep_actions - self.action_mean) / self.action_std

        if do_hflip and ep_actions.shape[1] >= 6:
            ep_actions = ep_actions.copy()
            ep_actions[:, 0] = -ep_actions[:, 0]
            ep_actions[:, 1] = -ep_actions[:, 1]
            ep_actions[:, 3] = -ep_actions[:, 3]
            ep_actions[:, 4] = -ep_actions[:, 4]

        T, D = ep_actions.shape
        K = self.action_chunk_size
        if T < K:
            ep_actions = np.concatenate([ep_actions, np.tile(ep_actions[-1:], (K - T, 1))], axis=0)
            T = K
        padded = np.concatenate([ep_actions, np.tile(ep_actions[-1:], (K - 1, 1))], axis=0)
        chunks = np.lib.stride_tricks.sliding_window_view(padded, (K, D))[:, 0, :, :]
        action_chunk = chunks[frame_idx].copy()

        if not self.use_processor:
            images = torch.stack([self.image_transform(img) for img in images])

        return {
            "image": images,
            "instruction": instruction,
            "state": torch.from_numpy(norm_state).float(),
            "action": torch.from_numpy(norm_action).float(),
            "action_chunk": torch.from_numpy(action_chunk).float(),
            "raw_action": torch.from_numpy(raw_action).float(),
            "frame_index": ep["raw_frame_indices"][frame_idx],
            "episode_index": ep["episode_index"],
            "dataset_name": self.dataset_name,
        }


# ---------------------------------------------------------------------------
# Collate 函数
# ---------------------------------------------------------------------------
def collate_fn(batch: List[Dict[str, Any]]) -> Dict[str, Any]:
    images = [b["image"] for b in batch]

    if isinstance(images[0], list):
        max_cam = max(len(imgs) for imgs in images)
        if max_cam > 1:
            for i in range(len(images)):
                if len(images[i]) < max_cam:
                    images[i] = images[i] + [images[i][0]] * (max_cam - len(images[i]))
    elif isinstance(images[0], torch.Tensor):
        max_cam = max(imgs.shape[0] for imgs in images)
        if max_cam > 1:
            for i in range(len(images)):
                if images[i].shape[0] < max_cam:
                    rep = images[i][0:1].repeat(max_cam - images[i].shape[0], 1, 1, 1)
                    images[i] = torch.cat([images[i], rep], dim=0)
        images = torch.stack(images)
        B, N, C, H, W = images.shape
        images = images.view(B * N, C, H, W)

    states = torch.stack([b["state"] for b in batch])
    actions = torch.stack([b["action"] for b in batch])
    action_chunks = torch.stack([b["action_chunk"] for b in batch])
    raw_actions = torch.stack([b["raw_action"] for b in batch])

    return {
        "image": images,
        "instruction": [b["instruction"] for b in batch],
        "state": states,
        "action": actions,
        "action_chunk": action_chunks,
        "raw_action": raw_actions,
        "frame_index": [b["frame_index"] for b in batch],
        "episode_index": [b["episode_index"] for b in batch],
        "dataset_name": [b["dataset_name"] for b in batch],
    }


# ---------------------------------------------------------------------------
# DataLoader 构建
# ---------------------------------------------------------------------------
def build_dataloader(config: dict) -> Tuple[DataLoader, Optional[DataLoader]]:
    """从 config 构建 (train_loader, val_loader)。"""
    data_cfg = config.get("data", {})
    model_cfg = config.get("model", {})
    train_cfg = config.get("training", {})

    action_dim = data_cfg.get("action_dim", model_cfg.get("action_dim", 7))
    state_dim = data_cfg.get("state_dim", model_cfg.get("state_dim", 7))
    action_chunk_size = data_cfg.get("action_chunk_size", model_cfg.get("action_chunk_size", 16))
    image_size = data_cfg.get("image_size", 224)
    use_processor = model_cfg.get("use_processor", True)
    val_ratio = data_cfg.get("val_ratio", 0.0)
    augment = data_cfg.get("domain_randomization", True)
    num_workers = data_cfg.get("num_workers", 1)
    pin_memory = data_cfg.get("pin_memory", True)

    dataset_configs = data_cfg.get("datasets", [])
    if not dataset_configs:
        dataset_configs = [{
            "name": "single",
            "dir": config.get("data_dir", ""),
            "weight": 1.0,
        }]

    action_adapter = data_cfg.get("action_adapter", None)
    state_adapter = data_cfg.get("state_adapter", None)

    train_dataset_list: List[LeRobotDataset] = []
    val_dataset_list: List[LeRobotDataset] = []

    for ds_cfg in dataset_configs:
        ds_dir = ds_cfg["dir"]
        ds_name = ds_cfg.get("name", Path(ds_dir).name)
        if not ds_dir or not Path(ds_dir).exists():
            logger.warning("[build_dataloader] dataset dir not found: %s", ds_dir)
            continue

        all_episodes = _collect_episode_indices(ds_dir)
        if not all_episodes:
            logger.warning("[build_dataloader] no episodes in %s", ds_dir)
            continue

        if val_ratio > 0 and len(all_episodes) > 1:
            n_val = max(1, int(len(all_episodes) * val_ratio))
            split_rng = random.Random(config.get("seed", 42))
            val_ids = set(split_rng.sample(sorted(all_episodes), n_val))
            train_ids = all_episodes - val_ids
        else:
            val_ids = None
            train_ids = all_episodes

        common_kwargs = {
            "data_dir": ds_dir,
            "dataset_name": ds_name,
            "action_dim": action_dim,
            "state_dim": state_dim,
            "action_chunk_size": action_chunk_size,
            "image_size": image_size,
            "action_normalize": data_cfg.get("action_normalize", True),
            "state_normalize": data_cfg.get("state_normalize", True),
            "frame_sampling": data_cfg.get("frame_sampling", "all"),
            "max_frames_per_episode": data_cfg.get("max_frames_per_episode", None),
            "use_processor": use_processor,
            "max_cache_episodes": data_cfg.get("max_cache_episodes", 200),
            "camera_keys": ds_cfg.get("camera_keys", data_cfg.get("camera_keys", ["observation.images.image_0"])),
            "dummy_image_on_missing_video": data_cfg.get("dummy_image_on_missing_video", True),
            "augment": augment,
            "action_adapter": action_adapter,
            "state_adapter": state_adapter,
        }

        train_ds = LeRobotDataset(val_episode_ids=train_ids, skip_stats_compute=False, **common_kwargs)
        train_dataset_list.append(train_ds)

        if val_ids:
            val_ds = LeRobotDataset(val_episode_ids=val_ids, skip_stats_compute=True,
                                    **{**common_kwargs, "augment": False})
            val_ds.set_stats(
                train_ds.action_mean, train_ds.action_std,
                train_ds.state_mean, train_ds.state_std,
            )
            val_dataset_list.append(val_ds)

    if not train_dataset_list:
        raise RuntimeError("No training datasets could be loaded. Check your config.")

    train_concat = ConcatDataset(train_dataset_list)
    dataset_weights = [ds_cfg.get("weight", 1.0) for ds_cfg in dataset_configs]

    sample_weights = []
    for ds, w in zip(train_dataset_list, dataset_weights):
        if len(ds) == 0:
            continue
        sample_weights.extend([w / max(1, len(ds))] * len(ds))
    sample_weights = torch.tensor(sample_weights, dtype=torch.double)

    total_train_samples = len(sample_weights)
    samples_per_epoch = data_cfg.get("samples_per_epoch", total_train_samples)
    num_samples = min(samples_per_epoch, total_train_samples)
    logger.info("[build_dataloader] samples_per_epoch=%d (total=%d, capped=%s)",
                samples_per_epoch, total_train_samples,
                "yes" if num_samples < total_train_samples else "no")

    sampler = WeightedRandomSampler(
        weights=sample_weights,
        num_samples=num_samples,
        replacement=True,
    )

    train_loader = DataLoader(
        train_concat,
        batch_size=train_cfg.get("batch_size", 16),
        sampler=sampler,
        num_workers=num_workers,
        pin_memory=pin_memory,
        collate_fn=collate_fn,
        drop_last=True,
        persistent_workers=num_workers > 0,
    )

    val_loader = None
    if val_dataset_list:
        val_concat = ConcatDataset(val_dataset_list)
        val_loader = DataLoader(
            val_concat,
            batch_size=train_cfg.get("batch_size", 16),
            shuffle=False,
            num_workers=num_workers,
            pin_memory=pin_memory,
            collate_fn=collate_fn,
            drop_last=False,
            persistent_workers=num_workers > 0,
        )

    logger.info("[build_dataloader] train datasets:")
    for ds in train_dataset_list:
        logger.info("  %s: %d samples", ds.dataset_name, len(ds))
    if val_loader:
        logger.info("[build_dataloader] val datasets:")
        for ds in val_dataset_list:
            logger.info("  %s: %d samples", ds.dataset_name, len(ds))

    _set_global_action_stats(train_concat, train_dataset_list, logger)

    return train_loader, val_loader


def _set_global_action_stats(
    train_concat: ConcatDataset,
    dataset_list: List["LeRobotDataset"],
    logger,
):
    """per-dataset 归一化 + 虚拟全局统计 fallback。"""
    total_samples = 0
    action_dim = dataset_list[0].action_dim
    weighted_mean = np.zeros(action_dim, dtype=np.float64)
    weighted_var = np.zeros_like(weighted_mean)

    for ds in dataset_list:
        n = len(ds)
        total_samples += n
        am = getattr(ds, "action_mean", None)
        as_ = getattr(ds, "action_std", None)
        if am is not None:
            weighted_mean += n * am.astype(np.float64)
        if as_ is not None:
            weighted_var += n * (as_.astype(np.float64) ** 2)

    if total_samples > 0:
        weighted_mean = (weighted_mean / total_samples).astype(np.float32)
        weighted_std = np.sqrt(weighted_var / total_samples).astype(np.float32)

        train_concat.action_mean = weighted_mean
        train_concat.action_std = weighted_std
        train_concat.state_mean = dataset_list[0].state_mean
        train_concat.state_std = dataset_list[0].state_std

        train_concat.per_dataset_stats = [
            {
                "name": ds.dataset_name,
                "action_mean": getattr(ds, "action_mean", None),
                "action_std": getattr(ds, "action_std", None),
                "state_mean": getattr(ds, "state_mean", None),
                "state_std": getattr(ds, "state_std", None),
            }
            for ds in dataset_list
        ]

        logger.info("[action_stats] per-dataset normalization (no global mixing):")
        for st in train_concat.per_dataset_stats:
            logger.info("  %s: mean=%s std=%s", st["name"], st["action_mean"], st["action_std"])
        logger.info("[action_stats] fallback global mean=%s std=%s (for model registration only)",
                    weighted_mean, weighted_std)
    else:
        logger.warning("[action_stats] no action stats available, using default (0, 1)")


def _collect_episode_indices(data_dir: str) -> set:
    """收集数据集里所有 episode_index (v2.0 + v3.0)。"""
    indices = set()
    base = Path(data_dir)
    patterns = [
        "data/chunk-*/episode_*.parquet",    # v2.0
        "data/chunk-*/file-*.parquet",       # v3.0
        "data__chunk-*__file-*.parquet",
        "*/data__chunk-*__file-*.parquet",
    ]
    files: List[Path] = []
    for pat in patterns:
        files.extend(base.glob(pat))
    for pf in set(files):
        try:
            df = pd.read_parquet(pf)
            if "episode_index" in df.columns:
                indices.update(df["episode_index"].unique().tolist())
            else:
                indices.add(0)
        except Exception:
            continue
    return indices


# ---------------------------------------------------------------------------
# 便捷测试
# ---------------------------------------------------------------------------
if __name__ == "__main__":
    import logging
    logging.basicConfig(level=logging.INFO)

    test_cfg = {
        "data_dir": "G:/LingArm/datasets/pretrain_pool/droid/sample_download",
        "model": {"use_processor": True},
        "training": {"batch_size": 2},
        "data": {
            "num_workers": 0,
            "image_size": 224,
            "action_chunk_size": 16,
            "action_dim": 7,
            "state_dim": 7,
            "action_normalize": True,
            "state_normalize": True,
            "domain_randomization": True,
            "datasets": [
                {"name": "droid", "dir": "G:/LingArm/datasets/pretrain_pool/droid/sample_download", "weight": 1.0},
            ],
        },
    }

    train_loader, val_loader = build_dataloader(test_cfg)
    batch = next(iter(train_loader))
    print("Batch keys:", batch.keys())
    print("image type:", type(batch["image"]))
    print("state shape:", batch["state"].shape)
    print("action shape:", batch["action"].shape)
    print("action_chunk shape:", batch["action_chunk"].shape)
    print("raw_action shape:", batch["raw_action"].shape)
    print("Instructions:", batch["instruction"][:2])
    print("Datasets:", batch["dataset_name"][:5])