Skip to content

feat: NVDEC GPU decode + parallel episode splitting + RoboMME training config#88

Open
Barristen wants to merge 3 commits into
Robbyant:mainfrom
Barristen:feat/gpu-nvdec-parallel-extract
Open

feat: NVDEC GPU decode + parallel episode splitting + RoboMME training config#88
Barristen wants to merge 3 commits into
Robbyant:mainfrom
Barristen:feat/gpu-nvdec-parallel-extract

Conversation

@Barristen
Copy link
Copy Markdown

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

Chenyi and others added 3 commits May 23, 2026 23:09
…raining 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 <noreply@anthropic.com>
- 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 <noreply@anthropic.com>
1-GPU, 5-step config pointing to HPC3 paths for testing the memory mechanism.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant