WSA: A 3D-Centric World-Spatial-Action Model
for Generalizable Robot Control

An embodied foundation model that introduces World-Spatial-Action modeling to unify instruction-aligned 2D visual planning, action-conditioned 3D world modeling, and 3D-aware action generation.

• [2026-07-06]: 🚀 Released WSA-Large code, weights, benchmark results, and training/evaluation workflows for the 6B Wan2.2-based WSA model.
• [2026-05-31]: 🎉 Released the WSA-Base training, evaluation, and inference code.
• [2026-05-31]: 🤗 Released the WSA-Base Hugging Face model collection, including Base, RoboTwin, and LIBERO checkpoints.
• [2026-05-18]: 🏆 WSA ranked 4th out of 100+ teams on the RoboChallenge CVPR 2026 leaderboard with fully open-source weights and code. (Team: MagicBot)

• Provide RoboTwin, LIBERO, and real-world robot example inference workflows.
• Release WSA policy code and fine-tuning scripts.
• Release WSA pretraining scripts.
• Release the WSA-Large code for the 6B WSA model based on the Wan2.2 model backbone.
• Release WSA-Large weights and benchmark results.
• Release the arXiv paper and citation.
• Provide the training and evaluation codes on RoboChallenge2.0.
• Release WSA1.5, our next-generation foundation model, pretrained on larger and more diverse datasets with memory capabilities for hard and long-horizon robotic tasks.

• Framework
• Repository Layout
• Installation
• Model Zoo
• Choosing a Model
• Inference
• Training
  • RoboTwin Fine-tuning
  • Fine-tuning example
  • WSA-Large Fine-tuning
  • Multi-Dataset Pretraining
• Acknowledgments
• Citation

WSA is a World-Spatial-Action (WSA) embodied foundation model for generalizable robot control. It learns a shared 2D-3D latent space that connects instruction-aligned visual planning, action-conditioned 3D world prediction, and 3D-aware action generation.

• Unified World-Spatial-Action Modeling: WSA unifies semantic understanding, 3D world modeling, and physical execution.
• Bidirectional 3D Causality: WSA learns both action-conditioned scene dynamics and 3D inverse dynamics.
• Mixture-of-Transformers: WSA coordinates 2D planning, 3D prediction, and 3D action generation with shared dependency rules.
• Data-Efficient Pretraining: Pretraining on 6,000 demonstration hours yields strong simulation and real-world manipulation performance.
• Superior Performance: State-of-the-art results across simulation and real-world robot manipulation tasks, achieved by our open-source model.

🤖 Result on RoboTwin 2.0 randomized setting, averaged over 50 simulated aloha manipulation tasks:

assets/                  README figures and paper assets
configs/                 data sampling and weight-rule configs
evaluation/
  RoboTwin/              RoboTwin evaluation entrypoints
  Libero/                LIBERO evaluation and websocket serving helpers
  Real_Piper_Example/    Piper real-robot serving/client example
  Real_Lift2_Example/    Lift2 real-robot serving/client example
launch/
  wsa_base_*.sh          WSA-B pretraining and fine-tuning scripts
  wsa_large_*.sh         WSA-L pretraining and fine-tuning scripts
  supported_methods/     RoboTwin fine-tuning scripts for comparison methods
src/lerobot/             LeRobot-based training, dataset, and policy code
third_party/             Git submodules for external projects
tools/                   support scripts used by training workflows

The main development environment uses Python 3.10, CUDA 12.8, and PyTorch 2.7.1.

git clone https://github.com/zaleni/WSA.git
cd WSA

conda create -y -n wsa python=3.10
conda activate wsa

conda install -c conda-forge ffmpeg=7.1.1 svt-av1 -y

pip install torch==2.7.1 torchvision==0.22.1 torchaudio==2.7.1 \
  --index-url https://download.pytorch.org/whl/cu128

pip install torchcodec numpy scipy transformers==4.57.1 mediapy loguru pytest omegaconf h5py rich
pip install -e .

WSA uses a patched Qwen3-VL implementation for cached inference. After installing transformers==4.57.1, copy the replacement model files into the installed package:

TRANSFORMERS_DIR=${CONDA_PREFIX}/lib/python3.10/site-packages/transformers/
cp -r src/lerobot/policies/WSA_Base/transformers_replace/models ${TRANSFORMERS_DIR}

RoboTwin 2.0 and LIBERO evaluation also require their official codebases. These dependencies are included as Git submodules under third_party/. To initialize them, run:

git submodule update --init --recursive

For real-robot serving and websocket evaluation:

pip install tyro matplotlib mediapy websockets msgpack

All released models are available in the WSA Hugging Face collection.

For action evaluation with the released model, use DISABLE_DA3_TEACHER_FOR_EVAL=true.

Use the task-specific checkpoint for evaluation. The plain WSA-Base and WSA-Large checkpoints are pretrained policies intended for downstream fine-tuning, while the *-RoboTwin and *-LIBERO checkpoints are ready for their corresponding benchmark adapters.

• RoboTwin: evaluation/RoboTwin/README.md
• LIBERO: evaluation/Libero/README.md
• Real Piper example: evaluation/Real_Piper_Example/README.md
• Real Lift2 example: evaluation/Real_Lift2_Example/README.md

The real-robot examples split inference into a GPU policy server and a robot-side client. They are intended as reference integrations that can be adapted to your own hardware. The released checkpoints were evaluated on NVIDIA GeForce RTX 4090 GPUs.

Note: The WSA-Large evaluator loads its text encoder by default so that it can accept plain-text instructions. For repeated deployment, you can instead precompute the exact task prompts with tools/precompute_text_embeds.py and run with WSA_LARGE_LOAD_TEXT_ENCODER=false; the benchmark and real-robot READMEs show the supported entrypoints.

All WSA training scripts are under launch/. For fine-tuning, initialize from the matching released pretrained checkpoint: POLICY_INIT_PATH=zaleni/WSA-Base or POLICY_INIT_PATH=zaleni/WSA-Large. Training used 8x NVIDIA H200 GPUs.

launch/wsa_base_finetune_robotwin.sh discovers all LeRobot-v3 datasets under ROBOTWIN_ROOT and trains on them as a multi-dataset run.

Download the RoboTwin LeRobot-v3.0 dataset from Hugging Face and point ROBOTWIN_ROOT to the local download directory:

hf download hxma/RoboTwin-LeRobot-v3.0 \
  --repo-type dataset \
  --local-dir /path/to/robotwin_lerobot_v3.0

Compute external normalization statistics before training. The output path below matches the DATASET_EXTERNAL_STATS_ROOT=/path/to/norm_stats layout used by the training script:

ROBOTWIN_ROOT=/path/to/robotwin_lerobot_v3.0

find -L "${ROBOTWIN_ROOT}" -path "*/meta/info.json" -print \
  | while read -r info; do dirname "$(dirname "$info")"; done \
  | sort -u > robotwin_repo_ids.txt

python tools/compute_norm_stats_multi.py \
  --repo_id_file robotwin_repo_ids.txt \
  --action_mode delta \
  --chunk_size 50 \
  --num_workers 8 \
  --output_path /path/to/norm_stats/aloha/delta/stats.json

If you want to train with ACTION_TYPE=abs, compute stats with --action_mode abs and write to /path/to/norm_stats/aloha/abs/stats.json instead.

POLICY_INIT_PATH=zaleni/WSA-Base \
ROBOTWIN_ROOT=/path/to/robotwin_lerobot_v3.0 \
ACTION_TYPE=delta \
USE_EXTERNAL_STATS=true \
DATASET_EXTERNAL_STATS_ROOT=/path/to/norm_stats \
bash launch/wsa_base_finetune_robotwin.sh

Use this script for a single LeRobot-v3.0 dataset. It defaults to delta actions.

POLICY_INIT_PATH=zaleni/WSA-Base \
DATASET_REPO_ID=/path/to/lerobot_v3.0_dataset \
ACTION_TYPE=delta \
USE_EXTERNAL_STATS=true \
bash launch/wsa_base_finetune.sh

For delta-action training, compute normalization statistics first:

python tools/compute_norm_stats_single.py \
  --repo_id /path/to/lerobot_v3.0_dataset \
  --action_mode delta \
  --chunk_size 50 \
  --output_dir norm_stats

WSA-Large additionally uses Wan text conditioning. The training launchers default to LOAD_TEXT_ENCODER=false, which keeps the text encoder out of the training process but requires a prompt-embedding cache. For RoboTwin, discover the datasets and build the cache first:

python tools/discover_robotwin_repos.py \
  --robotwin-root /path/to/robotwin_lerobot_v3.0 \
  --output-file robotwin_repo_ids.txt \
  --require-three-cameras true

python tools/precompute_text_embeds.py \
  --repo-id-file robotwin_repo_ids.txt \
  --text-embedding-cache-dir outputs/WSA_Large/text_embeds/robotwin \
  --device cuda

Compute normalization statistics with the same action mode and horizon used by training:

python tools/compute_norm_stats_multi.py \
  --repo_id_file robotwin_repo_ids.txt \
  --action_mode delta \
  --chunk_size 32 \
  --num_workers 8 \
  --output_path /path/to/norm_stats/aloha/delta/stats.json

Then fine-tune from the complete released WSA-Large checkpoint. Setting SKIP_DIT_LOAD_FROM_PRETRAIN=true tells the launcher to use the ActionDiT and Future3DExpert weights already contained in that checkpoint instead of asking for separately prepared expert-backbone files:

POLICY_INIT_PATH=zaleni/WSA-Large \
ROBOTWIN_ROOT=/path/to/robotwin_lerobot_v3.0 \
ACTION_TYPE=delta \
NORMALIZATION_STATS_PATH=/path/to/norm_stats/aloha/delta/stats.json \
TEXT_EMBED_CACHE_DIR=outputs/WSA_Large/text_embeds/robotwin \
LOAD_TEXT_ENCODER=false \
SKIP_DIT_LOAD_FROM_PRETRAIN=true \
bash launch/wsa_large_finetune_robotwin.sh

For a simpler bring-up at the cost of loading the Wan text encoder during training, set LOAD_TEXT_ENCODER=true and omit TEXT_EMBED_CACHE_DIR.

launch/wsa_base_pretrain.sh can discover datasets from multiple roots: INTERNDATA_ROOT, ROBOTWIN_ROOT, ROBOCHALLENGE_ROOT, AGIBOT_ROOT, and EGODEX_LEROBOT_ROOT.

ROBOTWIN_ROOT=/path/to/robotwin_lerobot_v3 \
EGODEX_LEROBOT_ROOT=/path/to/egodex_lerobot_v3 \
DATASET_EXTERNAL_STATS_ROOT=/path/to/norm_stats \
WEIGHT_RULES_PATH=configs/weight_rules_wsa_base_pretrain.yaml \
bash launch/wsa_base_pretrain.sh

For WSA-Large multi-dataset pretraining, prepare per-embodiment stats and use the WSA-Large launch script. Unlike fine-tuning from the complete released checkpoint, pretraining requires the two expert-backbone initialization files. Generate them once before launching:

python tools/preprocess_expert_backbones.py \
  --expert both \
  --action-output checkpoints/wsa_large/ActionDiT_linear_interp_Wan22_alphascale_1024hdim.pt \
  --future-3d-output checkpoints/wsa_large/Future3DExpert_linear_interp_Wan22_alphascale_768hdim.pt \
  --action-dim 24 \
  --da3-num-views 3 \
  --future-3d-tokens-per-view 144 \
  --device cuda \
  --dtype bfloat16

ROBOTWIN_ROOT=/path/to/robotwin_lerobot_v3 \
EGODEX_LEROBOT_ROOT=/path/to/egodex_lerobot_v3 \
OUTPUT_STATS_ROOT=/path/to/norm_stats \
bash tools/wsa_large_compute_pretrain_norm_stats.sh

python tools/precompute_text_embeds.py \
  --repo-id-file outputs/WSA_Large/_stats_repo_id_files/chunk32/all_datasets.txt \
  --text-embedding-cache-dir outputs/WSA_Large/text_embeds/pretrain \
  --device cuda

ROBOTWIN_ROOT=/path/to/robotwin_lerobot_v3 \
EGODEX_LEROBOT_ROOT=/path/to/egodex_lerobot_v3 \
DATASET_EXTERNAL_STATS_ROOT=/path/to/norm_stats \
TEXT_EMBED_CACHE_DIR=outputs/WSA_Large/text_embeds/pretrain \
WEIGHT_RULES_PATH=configs/weight_rules_wsa_large_pretrain.yaml \
bash launch/wsa_large_pretrain.sh

Some other policies are also supported by this repository, training scripts are available in launch/supported_methods/:

• qwenaction_finetune.sh
• pi0_finetune.sh
• pi05_finetune.sh
• internvla_a1_3b_finetune.sh
• fastwam_finetune.sh

WSA builds on the excellent work of the LeRobot, RoboTwin, Qwen3-VL, Depth-Anything-3, InternVLA-A1, and FastWAM. Some adapted policy scripts are kept in this repository to make reproduction and ablation runs easier from the same codebase.

If you find WSA useful in your research, please cite our paper:

@misc{jiang2026wsa,
  title         = {WSA$_1$: a 3D-Centric World-Spatial-Action Model for Generalizable Robot Control},
  author        = {Jiahao Jiang and Jianing Zhang and Zhenhan Yin and Ruidong Chen and Sen Wang and Zhaoshu Yu and Pengpeng Zeng and Xiaofeng Cao and Xuanhan Wang and Jingkuan Song and Heng Tao Shen},
  year          = {2026},
  eprint        = {2607.03941},
  archivePrefix = {arXiv},
  primaryClass  = {cs.RO},
  url           = {https://arxiv.org/abs/2607.03941}
}