From 3d7f2e9113a3cd499d6487e8dc95160c6ec29452 Mon Sep 17 00:00:00 2001 From: Chenyi <15810362+chenyi-zi@user.noreply.gitee.com> Date: Sat, 23 May 2026 23:09:55 +0800 Subject: [PATCH 1/3] feat: add NVDEC GPU decode, parallel episode splitting, and RoboMME training config - extract_latents.py: add NVDEC hardware video decoding with CPU fallback via PyAV hwaccel=cuda option; add --ep_start/--ep_end args for multi-GPU parallel episode range splitting across 8 GPUs - wan_va/configs/va_robomme_train_cfg.py: new training config for RoboMME dataset with 7-dim eef action space, normalization stats, and czi447 server paths - wan_va/configs/__init__.py: register robomme_train config in VA_CONFIGS dict Co-Authored-By: Claude Sonnet 4.6 --- extract_latents.py | 297 +++++++++++++++++++++++++ wan_va/configs/__init__.py | 2 + wan_va/configs/va_robomme_train_cfg.py | 67 ++++++ 3 files changed, 366 insertions(+) create mode 100644 extract_latents.py create mode 100644 wan_va/configs/va_robomme_train_cfg.py diff --git a/extract_latents.py b/extract_latents.py new file mode 100644 index 00000000..d2a66b35 --- /dev/null +++ b/extract_latents.py @@ -0,0 +1,297 @@ +""" +Extract VAE latents from LIBERO lerobot dataset for LingBot-VA training. + +Usage: + python extract_latents.py \ + --dataset_path /hpc2hdd/home/mpeng885/Datasets_raw/libero/libero_10_no_noops_1.0.0_lerobot \ + --model_path /hpc2hdd/home/mpeng885/chenyizi/lingbot-va-main/model \ + --obs_cam_keys observation.images.image observation.images.wrist_image \ + --target_fps 5 --height 128 --width 128 --device cuda:0 +""" + +import argparse +import json +from pathlib import Path + +import av +import numpy as np +import torch +import torch.nn.functional as F +from diffusers import AutoencoderKLWan +from einops import rearrange +from tqdm import tqdm +from transformers import T5TokenizerFast, UMT5EncoderModel + + +# --------------------------------------------------------------------------- +# Video reading +# --------------------------------------------------------------------------- + +def read_video_frames_av(video_path: Path, frame_ids: list, gpu_id: int = 0) -> list: + """Read specific frame indices from a video file using PyAV. + Tries NVDEC hardware decoding first, falls back to CPU on failure.""" + frame_set = set(frame_ids) + max_id = max(frame_ids) + buf = {} + + def _decode(container): + stream = container.streams.video[0] + stream.codec_context.skip_frame = 'DEFAULT' + for cur_idx, frame in enumerate(container.decode(stream)): + if cur_idx in frame_set: + buf[cur_idx] = frame.to_ndarray(format='rgb24') + if cur_idx > max_id: + break + + try: + with av.open(str(video_path), options={ + 'hwaccel': 'cuda', + 'hwaccel_device': str(gpu_id), + }) as container: + _decode(container) + except Exception: + buf.clear() + with av.open(str(video_path)) as container: + _decode(container) + + return [buf[fid] for fid in frame_ids if fid in buf] + + +def sample_frame_ids(start: int, end: int, ori_fps: int, target_fps: int) -> list: + """ + Sample frame indices at target_fps from [start, end). + Trims to largest (4n+1) count for VAE temporal alignment. + """ + stride = max(1, round(ori_fps / target_fps)) + ids = list(range(start, end, stride)) + n = (len(ids) - 1) // 4 + ids = ids[:4 * n + 1] + return ids if len(ids) >= 1 else [start] + + +def frames_to_tensor(frames: list, height: int, width: int) -> torch.Tensor: + """Convert list of HxWxC uint8 frames to [3, F, H, W] float32 in [-1, 1].""" + resized = [] + for f in frames: + t = torch.from_numpy(f).float() + t = t.permute(2, 0, 1).unsqueeze(0) # 1 C H W + t = F.interpolate(t, size=(height, width), mode='bilinear', align_corners=False) + resized.append(t.squeeze(0)) # C H W + tensor = torch.stack(resized, dim=1) # C F H W + return tensor / 127.5 - 1.0 + + +# --------------------------------------------------------------------------- +# Encoding helpers +# --------------------------------------------------------------------------- + +def encode_video(vae: AutoencoderKLWan, + video_tensor: torch.Tensor, + device: str): + """ + Encode [3, F, H, W] tensor -> (latent_flat [F'*H'*W', 48], F', H', W'). + Latent is normalized (z-scored) and stored as bfloat16. + Uses vae.encode() which handles full sequences correctly. + """ + x = video_tensor.unsqueeze(0).to(device).to(torch.bfloat16) # 1 3 F H W + with torch.no_grad(): + posterior = vae.encode(x).latent_dist + mu = posterior.mean # 1 48 F' H' W' + + mean = torch.tensor(vae.config.latents_mean, device=mu.device, dtype=torch.float32) + std = torch.tensor(vae.config.latents_std, device=mu.device, dtype=torch.float32) + mu_norm = ((mu.float() - mean.view(1, -1, 1, 1, 1)) / + std.view(1, -1, 1, 1, 1)).to(torch.bfloat16) + + mu_norm = mu_norm.squeeze(0).cpu() # 48 F' H' W' + _, Fl, Hl, Wl = mu_norm.shape + latent_flat = rearrange(mu_norm, 'c f h w -> (f h w) c') + return latent_flat, Fl, Hl, Wl + + +def encode_text(text_encoder, tokenizer, text: str, device: str, # noqa: E501 + max_length: int = 512) -> torch.Tensor: + """Encode text -> [L, D] bfloat16 tensor.""" + inputs = tokenizer(text, return_tensors='pt', padding='max_length', + max_length=max_length, truncation=True).to(device) + with torch.no_grad(): + emb = text_encoder(**inputs).last_hidden_state[0] # L D + return emb.cpu().to(torch.bfloat16) + + +# --------------------------------------------------------------------------- +# Main +# --------------------------------------------------------------------------- + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument('--dataset_path', required=True) + parser.add_argument('--model_path', required=True) + parser.add_argument('--obs_cam_keys', nargs='+', + default=['observation.images.image', + 'observation.images.wrist_image']) + parser.add_argument('--target_fps', type=int, default=5) + parser.add_argument('--height', type=int, default=128) + parser.add_argument('--width', type=int, default=128) + parser.add_argument('--device', default='cuda:0') + parser.add_argument('--ep_start', type=int, default=0) + parser.add_argument('--ep_end', type=int, default=-1) + args = parser.parse_args() + + dataset_path = Path(args.dataset_path) + model_path = Path(args.model_path) + device = args.device + gpu_id = int(device.split(':')[1]) if ':' in device else 0 + + # ---- Load models -------------------------------------------------------- + print('Loading VAE...') + vae = AutoencoderKLWan.from_pretrained( + str(model_path / 'vae'), torch_dtype=torch.bfloat16).to(device) + + print('Loading text encoder...') + text_encoder = UMT5EncoderModel.from_pretrained( + str(model_path / 'text_encoder'), torch_dtype=torch.bfloat16).to(device) + tokenizer = T5TokenizerFast.from_pretrained(str(model_path / 'tokenizer')) + + # ---- Empty embedding ---------------------------------------------------- + empty_emb_path = dataset_path / 'empty_emb.pt' + if not empty_emb_path.exists(): + print('Generating empty text embedding...') + empty_emb = encode_text(text_encoder, tokenizer, '', device) + torch.save(empty_emb, empty_emb_path) + print(f'Saved empty_emb.pt shape={empty_emb.shape}') + else: + print('empty_emb.pt already exists.') + + # ---- Load / update episodes.jsonl --------------------------------------- + episodes_file = dataset_path / 'meta' / 'episodes.jsonl' + episodes = [] + with open(episodes_file) as f: + for line in f: + line = line.strip() + if line: + episodes.append(json.loads(line)) + + changed = False + for ep in episodes: + if 'action_config' not in ep: + task_text = ep['tasks'][0] if ep['tasks'] else 'robot manipulation' + ep['action_config'] = [{ + 'start_frame': 0, + 'end_frame': ep['length'], + 'action_text': task_text, + }] + changed = True + if changed: + with open(episodes_file, 'w') as f: + for ep in episodes: + f.write(json.dumps(ep) + '\n') + print(f'Updated episodes.jsonl with action_config for {len(episodes)} episodes') + else: + print('episodes.jsonl already has action_config.') + + # ---- Dataset info ------------------------------------------------------- + with open(dataset_path / 'meta' / 'info.json') as f: + info = json.load(f) + ori_fps = info.get('fps', 20) + chunks_size = info.get('chunks_size', 1000) + print(f'Dataset: {len(episodes)} eps, ori_fps={ori_fps}, chunks_size={chunks_size}') + + # ---- Pre-encode text prompts ------------------------------------------- + all_texts = {acfg['action_text'] + for ep in episodes + for acfg in ep.get('action_config', [])} + print(f'Pre-encoding {len(all_texts)} unique prompts...') + text_cache = {} + for text in tqdm(sorted(all_texts), desc='text enc'): + text_cache[text] = encode_text(text_encoder, tokenizer, text, device) + + # ---- Extract latents ---------------------------------------------------- + print('Extracting VAE latents...') + skipped = processed = errors = 0 + + ep_end = args.ep_end if args.ep_end >= 0 else len(episodes) + episodes = episodes[args.ep_start:ep_end] + + for ep in tqdm(episodes, desc='episodes'): + ep_idx = ep['episode_index'] + ep_chunk = ep_idx // chunks_size + + for acfg in ep.get('action_config', []): + start_frame = acfg['start_frame'] + end_frame = acfg['end_frame'] + action_text = acfg['action_text'] + text_emb = text_cache[action_text] + + frame_ids = sample_frame_ids(start_frame, end_frame, ori_fps, args.target_fps) + if len(frame_ids) < 5: + print(f' Skip ep {ep_idx}: too few frames ({len(frame_ids)})') + continue + + # Check if all views already done + if all( + (dataset_path / 'latents' / f'chunk-{ep_chunk:03d}' / ck / + f'episode_{ep_idx:06d}_{start_frame}_{end_frame}.pth').exists() + for ck in args.obs_cam_keys + ): + skipped += 1 + continue + + # Encode each camera view + for cam_key in args.obs_cam_keys: + out_dir = (dataset_path / 'latents' / f'chunk-{ep_chunk:03d}' / cam_key) + out_file = out_dir / f'episode_{ep_idx:06d}_{start_frame}_{end_frame}.pth' + if out_file.exists(): + continue + + video_path = (dataset_path / 'videos' / f'chunk-{ep_chunk:03d}' / + cam_key / f'episode_{ep_idx:06d}.mp4') + if not video_path.exists(): + print(f' Warning: missing {video_path}') + errors += 1 + continue + + try: + frames = read_video_frames_av(video_path, frame_ids, gpu_id) + except Exception as e: + print(f' Error reading ep {ep_idx} {cam_key}: {e}') + errors += 1 + continue + + if len(frames) < 5: + print(f' Skip ep {ep_idx} {cam_key}: decoded only {len(frames)} frames') + continue + + # Trim to (4n+1) + n = (len(frames) - 1) // 4 + frames = frames[:4 * n + 1] + used_ids = frame_ids[:len(frames)] + + video_tensor = frames_to_tensor(frames, args.height, args.width) + latent_flat, Fl, Hl, Wl = encode_video(vae, video_tensor, device) + + out_dir.mkdir(parents=True, exist_ok=True) + torch.save({ + 'latent': latent_flat, # [F'*H'*W', 48] bfloat16 + 'latent_num_frames': Fl, + 'latent_height': Hl, + 'latent_width': Wl, + 'video_num_frames': len(frames), + 'video_height': args.height, + 'video_width': args.width, + 'text_emb': text_emb, # [L, D] bfloat16 + 'text': action_text, + 'frame_ids': used_ids, + 'start_frame': start_frame, + 'end_frame': end_frame, + 'fps': args.target_fps, + 'ori_fps': ori_fps, + }, out_file) + + processed += 1 + + print(f'\nDone! processed={processed} skipped={skipped} errors={errors}') + + +if __name__ == '__main__': + main() diff --git a/wan_va/configs/__init__.py b/wan_va/configs/__init__.py index fe3791d3..54484a1b 100644 --- a/wan_va/configs/__init__.py +++ b/wan_va/configs/__init__.py @@ -10,6 +10,7 @@ from .va_libero_cfg import va_libero_cfg from .va_libero_train_cfg import va_libero_train_cfg from .va_libero_i2va import va_libero_i2va_cfg +from .va_robomme_train_cfg import va_robomme_train_cfg VA_CONFIGS = { 'robotwin': va_robotwin_cfg, @@ -23,4 +24,5 @@ 'libero': va_libero_cfg, 'libero_train': va_libero_train_cfg, 'libero_i2av': va_libero_i2va_cfg, + 'robomme_train': va_robomme_train_cfg, } \ No newline at end of file diff --git a/wan_va/configs/va_robomme_train_cfg.py b/wan_va/configs/va_robomme_train_cfg.py new file mode 100644 index 00000000..0f307d91 --- /dev/null +++ b/wan_va/configs/va_robomme_train_cfg.py @@ -0,0 +1,67 @@ +# Copyright 2024-2025 The Robbyant Team Authors. All rights reserved. +from easydict import EasyDict +from .va_libero_cfg import va_libero_cfg + +va_robomme_train_cfg = EasyDict(__name__="Config: VA RoboMME train") +va_robomme_train_cfg.update(va_libero_cfg) + +# === Paths === +va_robomme_train_cfg.wan22_pretrained_model_name_or_path = "/hpc2hdd/home/czi447/lingbot-va-model" +va_robomme_train_cfg.dataset_path = "/hpc2hdd/home/czi447/robomme_lerobot" +va_robomme_train_cfg.empty_emb_path = "/hpc2hdd/home/czi447/robomme_lerobot/empty_emb.pt" +va_robomme_train_cfg.save_root = "/hpc2hdd/home/czi447/output/robomme_train" + +# === Camera keys === +va_robomme_train_cfg.obs_cam_keys = [ + "observation.images.image", + "observation.images.wrist_image", +] + +# === Action space (7-dim eef action, same as LIBERO) === +va_robomme_train_cfg.action_dim = 30 +va_robomme_train_cfg.used_action_channel_ids = list(range(0, 7)) +inverse_used_action_channel_ids = [7] * 30 # default to padding index +for i, j in enumerate(range(7)): + inverse_used_action_channel_ids[j] = i +va_robomme_train_cfg.inverse_used_action_channel_ids = inverse_used_action_channel_ids + +# === Normalization stats (from robomme assets/norm_stats.json, actions q01/q99) === +va_robomme_train_cfg.norm_stat = { + "q01": [ + -0.18040079298019407, + -0.3191569724082947, + -0.13117732753753664, + -0.5344759382843971, + -0.12701646758317942, + -0.2923886525869369, + -0.4027193263053894, + ] + [0.] * 23, + "q99": [ + 0.19587624197006226, + 0.3774391655921935, + 0.11749427890777586, + 0.4144394028544425, + 0.12551680266857157, + 0.5218227294564248, + 0.4166772034645081, + ] + [0.] * 23, +} + +# === Logging === +va_robomme_train_cfg.enable_wandb = False + +# === Data loading === +va_robomme_train_cfg.load_worker = 8 +va_robomme_train_cfg.save_interval = 200 +va_robomme_train_cfg.gc_interval = 50 +va_robomme_train_cfg.cfg_prob = 0.1 + +# === Training parameters === +va_robomme_train_cfg.learning_rate = 1e-5 +va_robomme_train_cfg.beta1 = 0.9 +va_robomme_train_cfg.beta2 = 0.95 +va_robomme_train_cfg.weight_decay = 1e-1 +va_robomme_train_cfg.warmup_steps = 10 +va_robomme_train_cfg.batch_size = 1 +va_robomme_train_cfg.gradient_accumulation_steps = 10 +va_robomme_train_cfg.num_steps = 20000 From e3b0e65d0f960c4712a0e27932b439b3ae987a72 Mon Sep 17 00:00:00 2001 From: Chenyi <15810362+chenyi-zi@user.noreply.gitee.com> Date: Sun, 24 May 2026 12:27:12 +0800 Subject: [PATCH 2/3] feat: add memory mechanism to LingBot-VA (Route 1) - Dataset: extract memory_latents (first K frames spatially mean-pooled) from episode latent - Model: add memory_proj linear, inject memory tokens into cross-attention alongside text - Train: pass memory_latents from batch through _prepare_input_dict into model Co-Authored-By: Claude Sonnet 4.6 --- wan_va/dataset/lerobot_latent_dataset.py | 11 +++++++++++ wan_va/modules/model.py | 7 +++++++ wan_va/train.py | 2 ++ 3 files changed, 20 insertions(+) diff --git a/wan_va/dataset/lerobot_latent_dataset.py b/wan_va/dataset/lerobot_latent_dataset.py index b2860874..8003b2ce 100644 --- a/wan_va/dataset/lerobot_latent_dataset.py +++ b/wan_va/dataset/lerobot_latent_dataset.py @@ -306,6 +306,17 @@ def __getitem__(self, idx) -> dict: out_dict['actions'], out_dict['actions_mask'] = self._action_post_process(local_start_frame, local_end_frame, latent_frame_ids, ori_data_dict['action']) out_dict['latents'] = out_dict['latents'].permute(3, 0, 1, 2) + + # Memory tokens: spatial mean-pool first K frames of first camera's latent + first_key = self.used_video_keys[0] + lat = ori_data_dict[f"{first_key}.latent"] # [F*H*W, 48] + Fl = ori_data_dict[f"{first_key}.latent_num_frames"] + Hl = ori_data_dict[f"{first_key}.latent_height"] + Wl = ori_data_dict[f"{first_key}.latent_width"] + lat_3d = lat.reshape(Fl, Hl * Wl, 48) # [F, HW, 48] + mem_budget = min(16, Fl) + out_dict['memory_latents'] = lat_3d[:mem_budget].mean(dim=1) # [budget, 48] + return out_dict def __len__(self): diff --git a/wan_va/modules/model.py b/wan_va/modules/model.py index 25b45e57..cad1de70 100644 --- a/wan_va/modules/model.py +++ b/wan_va/modules/model.py @@ -625,6 +625,7 @@ def __init__(self, in_channels * patch_size[0] * patch_size[1] * patch_size[2], inner_dim) self.action_embedder = nn.Linear(action_dim, inner_dim) + self.memory_proj = nn.Linear(in_channels, inner_dim) self.condition_embedder = WanTimeTextImageEmbedding( dim=inner_dim, time_freq_dim=freq_dim, @@ -715,6 +716,12 @@ def forward_train(self, input_dict): text_hidden_states = text_hidden_states.flatten(0, 1)[None] + if 'memory_latents' in input_dict: + memory = input_dict['memory_latents'].to(latent_hidden_states.dtype) # B budget 48 + memory = self.memory_proj(memory) # B budget inner_dim + memory = memory.flatten(0, 1)[None] # 1 B*budget inner_dim + text_hidden_states = torch.cat([text_hidden_states, memory], dim=1) + condition_latent_hidden_states = self._input_embed(latent_dict['latent'], input_type='latent').flatten(0, 1)[None] condition_action_hidden_states = self._input_embed(action_dict['latent'], input_type='action').flatten(0, 1)[None] diff --git a/wan_va/train.py b/wan_va/train.py index ea3a7d5b..cf317333 100644 --- a/wan_va/train.py +++ b/wan_va/train.py @@ -245,6 +245,8 @@ def _prepare_input_dict(self, batch_dict): 'chunk_size': torch.randint(1, 5, (1,)).item(), 'window_size': torch.randint(4, 65, (1,)).item(), } + if 'memory_latents' in batch_dict: + input_dict['memory_latents'] = batch_dict['memory_latents'] return input_dict def convert_input_format(self, input_dict): From 79ac7317523e3562e78229548397d9ffc6cb5d96 Mon Sep 17 00:00:00 2001 From: Chenyi <15810362+chenyi-zi@user.noreply.gitee.com> Date: Sun, 24 May 2026 12:28:44 +0800 Subject: [PATCH 3/3] feat: add robomme_debug config for quick validation 1-GPU, 5-step config pointing to HPC3 paths for testing the memory mechanism. Co-Authored-By: Claude Sonnet 4.6 --- wan_va/configs/__init__.py | 2 ++ wan_va/configs/va_robomme_debug_cfg.py | 19 +++++++++++++++++++ 2 files changed, 21 insertions(+) create mode 100644 wan_va/configs/va_robomme_debug_cfg.py diff --git a/wan_va/configs/__init__.py b/wan_va/configs/__init__.py index 54484a1b..65c6e1c2 100644 --- a/wan_va/configs/__init__.py +++ b/wan_va/configs/__init__.py @@ -11,6 +11,7 @@ from .va_libero_train_cfg import va_libero_train_cfg from .va_libero_i2va import va_libero_i2va_cfg from .va_robomme_train_cfg import va_robomme_train_cfg +from .va_robomme_debug_cfg import va_robomme_debug_cfg VA_CONFIGS = { 'robotwin': va_robotwin_cfg, @@ -25,4 +26,5 @@ 'libero_train': va_libero_train_cfg, 'libero_i2av': va_libero_i2va_cfg, 'robomme_train': va_robomme_train_cfg, + 'robomme_debug': va_robomme_debug_cfg, } \ No newline at end of file diff --git a/wan_va/configs/va_robomme_debug_cfg.py b/wan_va/configs/va_robomme_debug_cfg.py new file mode 100644 index 00000000..c1161952 --- /dev/null +++ b/wan_va/configs/va_robomme_debug_cfg.py @@ -0,0 +1,19 @@ +# Copyright 2024-2025 The Robbyant Team Authors. All rights reserved. +from easydict import EasyDict +from .va_robomme_train_cfg import va_robomme_train_cfg + +va_robomme_debug_cfg = EasyDict(__name__="Config: VA RoboMME debug") +va_robomme_debug_cfg.update(va_robomme_train_cfg) + +# === Paths (HPC3) === +va_robomme_debug_cfg.wan22_pretrained_model_name_or_path = "/data/user/czi447/lingbot-va-model" +va_robomme_debug_cfg.dataset_path = "/data/user/czi447/robomme_lerobot" +va_robomme_debug_cfg.empty_emb_path = "/data/user/czi447/robomme_lerobot/empty_emb.pt" +va_robomme_debug_cfg.save_root = "/data/user/czi447/output/robomme_debug" + +# === Debug: minimal steps, no wandb === +va_robomme_debug_cfg.num_steps = 5 +va_robomme_debug_cfg.save_interval = 999999 +va_robomme_debug_cfg.gradient_accumulation_steps = 1 +va_robomme_debug_cfg.enable_wandb = False +va_robomme_debug_cfg.load_worker = 2