CoolFace
Datasetpublic

MangoGoes/libero4in1_wan2.2vae_latent_dataset

LIBERO 4in1 Wan2.2-VAE Latent Cache Pre-encoded latent tensors for LIBERO 4 suites under Cosmos Wan2.2-VAE. Skip on-the-fly VAE encoding during training — load this cache directly. Overview Pre-encoded latent cache for LIBERO 4in1 benchmark (libero_spatial, libero_object, libero_goal, libero_10 — 4 suites × 10 tasks, ~1700 episodes total). Each raw video frame is encoded once with Wan2.2-VAE, then saved as .pt tensors for direct loading during action-policy… See the full description on the dataset page: https://huggingface.co/datasets/MangoGoes/libero4in1_wan2.2vae_latent_dataset.

sourceHugging Faceupdated 1mo agoView on Hugging Face
0likes235downloads
Dataset Card

LIBERO 4in1 Wan2.2-VAE Latent Cache

Pre-encoded latent tensors for LIBERO 4 suites under Cosmos Wan2.2-VAE. Skip on-the-fly VAE encoding during training — load this cache directly.

Overview

Pre-encoded latent cache for LIBERO 4in1 benchmark (libero_spatial, libero_object, libero_goal, libero_10 — 4 suites × 10 tasks, ~1700 episodes total). Each raw video frame is encoded once with Wan2.2-VAE, then saved as .pt tensors for direct loading during action-policy training (action head / VLA / SFT).

Encoding once at ~5–10 GB/s inference throughput and reading .pt files at train time avoids repeated VAE calls, yielding a 5–10× training speedup.

Encoding Configuration

ItemValue
VAEWan2.2_VAE.pth (from Cosmos3 Edge project, channels=16, latent_dim=16)
Frame size256 × 256
Frame stackingMulti-frame window (--windowed); window length comes from dataset meta
Encoded dtypefloat16 (storage-efficient)
Output formatOne .pt per episode: {"video": Tensor[T,C,H,W], "state": Tensor[T, D_state], "action": Tensor[T, D_action]}
Manifestdataset_manifest*.json (indexed by taskid / episodeid)

Encoder script: tools/g0/build_cosmos_libero_latent_dataset.py

  • —Path: /disk/rl/psm_wma/tools/g0/build_cosmos_libero_latent_dataset.py
  • —Run as: single suite + multi-shard parallelism (3–5 shards depending on GPU memory)

Directory Structure

libero4in1_wan2.2vae_latent_dataset/
├── README.md
├── libero_spatial/
│   ├── dataset_manifest.json
│   ├── dataset_manifest_shard_0000.json
│   ├── dataset_manifest_shard_0001.json
│   ├── dataset_manifest_shard_0002.json
│   └── episodes/
│       ├── episode_000000.pt      # ~24 MB / episode
│       ├── episode_000001.pt
│       └── ...                    # 438 episodes
├── libero_object/                 # 460 episodes
├── libero_goal/                   # 434 episodes
└── libero_10/                     # 385 episodes

Dataset Size

SuiteTasksEpisodesSize
libero_spatial104389.94 GB
libero_object1046012.89 GB
libero_goal104349.76 GB
libero_101038520.60 GB
Total40171753.19 GB

Usage

1. huggingface_hub.snapshot_download (recommended)

python
import os
os.environ['HF_HUB_ENABLE_HF_TRANSFER'] = '1'
from huggingface_hub import snapshot_download

local_dir = snapshot_download(
    repo_id='MangoGoes/libero4in1_wan2.2vae_latent_dataset',
    repo_type='dataset',
    local_dir='/disk/rl/data/libero4in1_latent',
)
print(local_dir)
# → /disk/rl/data/libero4in1_latent

2. hf CLI

bash
hf download MangoGoes/libero4in1_wan2.2vae_latent_dataset \
  --repo-type dataset \
  --local-dir /disk/rl/data/libero4in1_latent

3. Custom torch.utils.data.Dataset

python
import json, torch
from pathlib import Path
from torch.utils.data import Dataset

class LiberoLatentDataset(Dataset):
    def __init__(self, cache_root, suite='libero_spatial'):
        self.root = Path(cache_root) / suite / 'episodes'
        with open(Path(cache_root) / suite / 'dataset_manifest.json') as f:
            self.manifest = json.load(f)
        self.episodes = sorted(self.root.glob('episode_*.pt'))

    def __len__(self):
        return len(self.episodes)

    def __getitem__(self, idx):
        return torch.load(self.episodes[idx], weights_only=True)

4. Single-file random sampling

python
import random, torch
from pathlib import Path

ep = random.choice(
    list(Path('/disk/rl/data/libero4in1_latent/libero_spatial/episodes').glob('*.pt'))
)
data = torch.load(ep, weights_only=True)
video  = data['video']    # [T, C, H, W] float16
state  = data['state']    # [T, D_state] float32
action = data['action']   # [T, D_action] float32

Related Resources

Mirror

A mirror exists at the original model-typed repo (kept for historical reference; HF does not allow changing repo_type post-creation):

  • —https://huggingface.co/MangoGoes/libero4in1wan2.2vaelatentcosmosstyle

Files are identical between the two repos; please prefer this dataset-typed repo for new downloads.

Re-encoding Command

bash
REPO_ROOT=/disk/rl/psm_wma
VAE_PATH="$REPO_ROOT/cosmos-framework/examples/checkpoints/wan22_vae/Wan2.2_VAE.pth"
OUTPUT_ROOT=/disk/rl/data/LIBERO_LeRobot_v3_cosmos_exact_window_shared_vae_v1

for shard in 0 1 2; do
  "$REPO_ROOT/cosmos-framework/.venv/bin/python" \
    "$REPO_ROOT/tools/g0/build_cosmos_libero_latent_dataset.py" \
    --dataset-root "/disk/rl/data/LIBERO_LeRobot_v3/libero_spatial" \
    --output-root "$OUTPUT_ROOT/libero_spatial" \
    --vae-path "$VAE_PATH" \
    --image-size 256 \
    --device cuda \
    --windowed \
    --episode-shard "$shard" \
    --num-shards 3
done