[Train] Support Fastwam training

lightx2v_train/configs/train/fastwam/libero_uncond_2cam224.yaml

@@ -0,0 +1,62 @@
+model:
+  name: wan_fastwam
+  running_dtype: bf16
+  model_path: /path/to/Wan-AI/Wan2.2-TI2V-5B
+
+data:
+  train: &libero_data
+    name: libero_fastwam_dataset
+    dataset_dirs:
+      # https://huggingface.co/datasets/yuanty/LIBERO-fastwam
+      - /path/to/LIBERO-fastwam/libero_spatial_no_noops_lerobot
+      - /path/to/LIBERO-fastwam/libero_object_no_noops_lerobot
+      - /path/to/LIBERO-fastwam/libero_goal_no_noops_lerobot
+      - /path/to/LIBERO-fastwam/libero_10_no_noops_lerobot
+    batch_size: 8
+    num_workers: 8
+  val:
+    <<: *libero_data
+
+training:
+  method: fastwam
+  output_dir: runs/fastwam_libero_full
+  # Equivalent to the official 10-epoch, global-batch-128 LIBERO schedule.
+  max_train_iters: 21700
+  gradient_accumulation_iters: 4
+  max_grad_norm: 1.0
+  save_every_iters: 2000
+  save_total_limit: 2
+  save_final: true
+  lr_scheduler: cosine
+  lr_warmup_iters: 1085
+  optimizer:
+    learning_rate: 1.0e-4
+    weight_decay: 1.0e-2
+    adam_beta1: 0.9
+    adam_beta2: 0.95
+    adam_epsilon: 1.0e-8
+
+evaluation:
+  eval_every_iters: 200
+  eval_num_inference_steps: 5
+  num_samples: 1
+  run_inference: true
+  save_video: true
+  save_preview: true
+  fps: 8
+  tiled: false
+  seed: 42
+
+logging:
+  train_log_every_iters: 10
+
+resume:
+  auto_resume: true
+  resume_ckpt_path: null
+
+distributed:
+  backend: nccl
+  zero1:
+    enabled: true
+  sequence_parallel:
+    enabled: false

lightx2v_train/lightx2v_train/data/__init__.py

@@ -1,13 +1,17 @@
 from lightx2v_train.utils.registry import build_data
 
 from .image_dataset import build_image_dataset
+from .libero.dataset import build_libero_fastwam_dataset
+from .preparation import prepare_data
 from .video_dataset import build_causal_forcing_lmdb_dataset, build_prompt_dataset, build_wan_t2v_cached_dataset, build_wan_t2v_video_dataset
 
 __all__ = [
     "build_data",
     "build_image_dataset",
+    "build_libero_fastwam_dataset",
     "build_prompt_dataset",
     "build_wan_t2v_video_dataset",
     "build_wan_t2v_cached_dataset",
     "build_causal_forcing_lmdb_dataset",
+    "prepare_data",
 ]

lightx2v_train/lightx2v_train/data/libero/__init__.py

@@ -0,0 +1,3 @@
+from .dataset import build_libero_fastwam_dataset
+
+__all__ = ["build_libero_fastwam_dataset"]

lightx2v_train/lightx2v_train/data/libero/dataset.py

@@ -0,0 +1,103 @@
+from copy import deepcopy
+from pathlib import Path
+
+import torch
+from loguru import logger
+from torch.utils.data import DataLoader, DistributedSampler
+
+from lightx2v_train.runtime.distributed import get_data_parallel_rank, get_data_parallel_world_size
+from lightx2v_train.utils.registry import DATA_REGISTER
+
+from .processor import FastWAMProcessor
+from .robot_video_dataset import RobotVideoDataset
+
+
+def _default_shape_meta():
+    return {
+        "images": [
+            {"key": "image", "raw_shape": [3, 512, 512], "shape": [3, 224, 224]},
+            {"key": "wrist_image", "raw_shape": [3, 512, 512], "shape": [3, 224, 224]},
+        ],
+        "action": [{"key": "default", "raw_shape": 7, "shape": 7}],
+        "state": [{"key": "default", "raw_shape": 8, "shape": 8}],
+    }
+
+
+class DatasetSliceRepeat(torch.utils.data.Dataset):
+    def __init__(self, dataset, max_samples=None, dataset_repeat=1):
+        self.dataset = dataset
+        self.base_len = len(dataset) if max_samples is None else min(int(max_samples), len(dataset))
+        self.dataset_repeat = max(1, int(dataset_repeat))
+        if self.base_len <= 0:
+            raise RuntimeError("LIBERO dataset is empty after applying max_samples.")
+
+    def __len__(self):
+        return self.base_len * self.dataset_repeat
+
+    def __getitem__(self, index):
+        return self.dataset[index % self.base_len]
+
+
+def _path(value):
+    return str(Path(value).expanduser().resolve())
+
+
+def _build_dataset(config, split):
+    shape_meta = deepcopy(config.get("shape_meta") or _default_shape_meta())
+    num_frames = int(config.get("num_frames", 33))
+    processor = FastWAMProcessor(shape_meta, num_frames)
+    dataset_dirs = config.get("dataset_dirs")
+    if isinstance(dataset_dirs, (str, Path)):
+        dataset_dirs = [dataset_dirs]
+
+    dataset = RobotVideoDataset(
+        dataset_dirs=[_path(item) for item in dataset_dirs],
+        shape_meta=shape_meta,
+        processor=processor,
+        text_embedding_cache_dir=_path(config["text_embedding_cache_dir"]),
+        pretrained_norm_stats=_path(config["pretrained_norm_stats"]),
+        num_frames=num_frames,
+        context_len=int(config.get("context_len", 128)),
+        val_set_proportion=float(config.get("val_set_proportion", 0.0)),
+        is_training_set=bool(config.get("is_training_set", split == "train")),
+        global_sample_stride=int(config.get("global_sample_stride", 1)),
+        action_video_freq_ratio=int(config.get("action_video_freq_ratio", 4)),
+        skip_padding_as_possible=bool(config.get("skip_padding_as_possible", False)),
+        max_padding_retry=int(config.get("max_padding_retry", 3)),
+        video_backend=config.get("video_backend"),
+    )
+    logger.info("[data] built LIBERO FastWAM {} dataset size={}", split, len(dataset))
+    return DatasetSliceRepeat(
+        dataset,
+        max_samples=config.get("max_samples"),
+        dataset_repeat=config.get("dataset_repeat", 1),
+    )
+
+
+@DATA_REGISTER("libero_fastwam_dataset")
+def build_libero_fastwam_dataset(data_config, train_or_val="train"):
+    dataset = _build_dataset(data_config, train_or_val)
+    if data_config.get("return_dataset", False):
+        return dataset
+
+    sampler = None
+    shuffle = bool(data_config.get("shuffle", train_or_val == "train"))
+    world_size = get_data_parallel_world_size()
+    if train_or_val == "train" and world_size > 1:
+        sampler = DistributedSampler(
+            dataset,
+            num_replicas=world_size,
+            rank=get_data_parallel_rank(),
+            shuffle=shuffle,
+            drop_last=bool(data_config.get("drop_last", False)),
+        )
+        shuffle = False
+    return DataLoader(
+        dataset,
+        batch_size=int(data_config.get("batch_size", 1)),
+        shuffle=shuffle,
+        sampler=sampler,
+        num_workers=int(data_config.get("num_workers", 4)),
+        pin_memory=bool(data_config.get("pin_memory", torch.cuda.is_available())),
+        drop_last=bool(data_config.get("drop_last", False)),
+    )

lightx2v_train/lightx2v_train/data/libero/lerobot_dataset.py

@@ -0,0 +1,187 @@
+import bisect
+import json
+from dataclasses import dataclass
+from functools import lru_cache
+from pathlib import Path
+
+import numpy as np
+import pyarrow as pa
+import pyarrow.parquet as pq
+import torch
+from torch.utils.data import Dataset
+
+from .video_decoder import decode_video_frames
+
+
+@dataclass(frozen=True)
+class Episode:
+    root: Path
+    index: int
+    length: int
+    data_path: str
+    video_path: str
+    chunks_size: int
+    tasks: dict[int, str]
+
+    def parquet_path(self) -> Path:
+        return self.root / self.data_path.format(
+            episode_chunk=self.index // self.chunks_size,
+            episode_index=self.index,
+        )
+
+    def camera_path(self, key: str) -> Path:
+        return self.root / self.video_path.format(
+            episode_chunk=self.index // self.chunks_size,
+            episode_index=self.index,
+            video_key=key,
+        )
+
+
+def _read_json(path: Path):
+    with path.open(encoding="utf-8") as handle:
+        return json.load(handle)
+
+
+def _read_jsonl(path: Path):
+    with path.open(encoding="utf-8") as handle:
+        return [json.loads(line) for line in handle if line.strip()]
+
+
+def _column_to_tensor(column: pa.ChunkedArray) -> torch.Tensor:
+    column = column.combine_chunks()
+    if pa.types.is_list(column.type) or pa.types.is_fixed_size_list(column.type):
+        array = np.asarray(column.to_pylist())
+    else:
+        array = column.to_numpy(zero_copy_only=False)
+    return torch.from_numpy(np.array(array, copy=True))
+
+
+class LiberoLeRobotDataset(Dataset):
+    """Minimal local reader for FastWAM's processed LIBERO LeRobot v2.1 data."""
+
+    def __init__(
+        self,
+        dataset_dirs,
+        image_keys,
+        state_key,
+        action_key,
+        num_frames,
+        global_sample_stride=1,
+        val_set_proportion=0.0,
+        is_training_set=True,
+        seed=42,
+        video_backend=None,
+    ):
+        if num_frames < 2:
+            raise ValueError(f"num_frames must be at least 2, got {num_frames}")
+        self.image_keys = list(image_keys)
+        self.state_key = state_key
+        self.action_key = action_key
+        self.num_frames = int(num_frames)
+        self.global_sample_stride = int(global_sample_stride)
+        self.video_backend = video_backend
+
+        episodes = []
+        fps_values = set()
+        for dataset_dir in dataset_dirs:
+            root = Path(dataset_dir).expanduser().resolve()
+            info = _read_json(root / "meta" / "info.json")
+            tasks = {int(item["task_index"]): item["task"] for item in _read_jsonl(root / "meta" / "tasks.jsonl")}
+            episode_meta = _read_jsonl(root / "meta" / "episodes.jsonl")
+            fps_values.add(int(info["fps"]))
+            self._validate_features(root, info["features"])
+
+            indices = list(range(len(episode_meta)))
+            if val_set_proportion >= 1e-6:
+                rng = np.random.default_rng(seed)
+                rng.shuffle(indices)
+                split = int(len(indices) * (1.0 - val_set_proportion))
+                indices = indices[:split] if is_training_set else indices[split:]
+
+            by_index = {int(item["episode_index"]): item for item in episode_meta}
+            for episode_index in indices:
+                item = by_index[episode_index]
+                episodes.append(
+                    Episode(
+                        root=root,
+                        index=episode_index,
+                        length=int(item["length"]),
+                        data_path=info["data_path"],
+                        video_path=info["video_path"],
+                        chunks_size=int(info["chunks_size"]),
+                        tasks=tasks,
+                    )
+                )
+
+        if not episodes:
+            raise RuntimeError("No LIBERO episodes were selected.")
+        if len(fps_values) != 1:
+            raise ValueError(f"All LIBERO datasets must have the same fps, got {sorted(fps_values)}")
+        self.fps = fps_values.pop()
+        self.episodes = episodes
+        self._episode_ends = np.cumsum([episode.length for episode in episodes]).tolist()
+
+    def _validate_features(self, root, features):
+        required = [*self.image_keys, self.state_key, self.action_key, "timestamp", "task_index"]
+        missing = [key for key in required if key not in features]
+        if missing:
+            raise ValueError(f"LIBERO dataset {root} is missing features: {missing}")
+        non_video = [key for key in self.image_keys if features[key].get("dtype") != "video"]
+        if non_video:
+            raise ValueError(f"LIBERO camera features must be videos, got: {non_video}")
+
+    def __len__(self):
+        return self._episode_ends[-1]
+
+    @lru_cache(maxsize=32)
+    def _load_episode(self, episode_position):
+        episode = self.episodes[episode_position]
+        table = pq.read_table(episode.parquet_path())
+        return {key: _column_to_tensor(table[key]) for key in table.column_names}
+
+    def _locate(self, index):
+        if index < 0:
+            index += len(self)
+        if index < 0 or index >= len(self):
+            raise IndexError(f"Index {index} out of bounds for dataset of size {len(self)}")
+        episode_position = bisect.bisect_right(self._episode_ends, index)
+        episode_start = 0 if episode_position == 0 else self._episode_ends[episode_position - 1]
+        return episode_position, index - episode_start
+
+    def _window(self, values, frame_index, length):
+        offsets = torch.arange(length, dtype=torch.long) * self.global_sample_stride
+        indices = frame_index + offsets
+        is_pad = indices >= values.shape[0]
+        indices.clamp_(max=values.shape[0] - 1)
+        return values[indices], is_pad, indices
+
+    def __getitem__(self, index):
+        episode_position, frame_index = self._locate(index)
+        episode = self.episodes[episode_position]
+        data = self._load_episode(episode_position)
+
+        state, state_is_pad, observation_indices = self._window(data[self.state_key], frame_index, self.num_frames)
+        action, action_is_pad, _ = self._window(data[self.action_key], frame_index, self.num_frames - 1)
+        timestamps = data["timestamp"][observation_indices].float().tolist()
+        tolerance = max(1e-4, 1.0 / self.fps - 1e-4)
+        images = {
+            key: decode_video_frames(
+                episode.camera_path(key),
+                timestamps,
+                tolerance_s=tolerance,
+                backend=self.video_backend,
+            )
+            for key in self.image_keys
+        }
+
+        task_index = int(data["task_index"][frame_index].item())
+        return {
+            "idx": int(index),
+            "task": episode.tasks[task_index],
+            "action": {"default": action.float()},
+            "state": {"default": state.float()},
+            "images": {key.removeprefix("observation.images."): value for key, value in images.items()},
+            "action_is_pad": action_is_pad,
+            "state_is_pad": state_is_pad,
+            "image_is_pad": state_is_pad.clone(),
+        }