The pretrained & post-trained checkpoint (simulation), inference server, and post-training toolkit are now available.
This repository releases the model, the inference server, and the post-training
toolkit, so you can load our pretrained weights and adapt
Real-robot demo highlights — click for the full video on the
project page.
- Load the pretrained
$N_0$ -TWAM checkpoint (n0_twam.models.utils.load_mot_checkpoint). - Post-train it on your own demonstrations (
n0_twam/train.py) — see POST_TRAINING.md. - Serve it over a websocket and get actions from observations (
n0_twam/n0_twam_server.py) — see DEPLOY.md. - Evaluate it closed-loop in the NeoSim vision–tactile benchmark — see Evaluate in NeoSim.
- Drive it closed-loop from your own robot or simulator (
example_client/closed_loop_client.py, numpy-only) — see DEPLOY.md.
| Component | Choice |
|---|---|
| Backbone | WAN2.2 TI2V-5B video diffusion transformer, restructured into a 3-expert MoT |
| Video VAE | Wan2.2 AutoencoderKLWan (z_dim=48, 4× temporal / 16× spatial); 256×256×129 → 48×33×16×16 latent |
| Text encoder | umT5-xxl (4096-d), frozen |
| Objective | Rectified-flow / flow-matching (FlowMatchScheduler), per-frame timesteps |
| Resolution / SNR shift | video & tactile snr_shift=5.0, action snr_shift=1.0 |
| Precision | bf16 parameters, fp32 reductions (FSDP2 MixedPrecisionPolicy) |
| Action space | 20-dim dual-arm end-effector, π0.5-style horizon delta |
See the Highlights section at the bottom for the modeling ideas.
Model overview: a Mixture-of-Transformers over joint video / tactile / action tokens.
n0-twam/
├── n0_twam/ # the model package
│ ├── models/ # MoT backbone
│ │ ├── mot.py # Mixture-of-Transformers (per-modality experts + shared attn)
│ │ ├── model.py # transformer blocks, tactile & action heads
│ │ └── utils.py # load_mot_checkpoint, VAE / text-encoder loaders
│ ├── configs/ # config registry (TWAM_CONFIGS)
│ │ ├── shared_config.py # shared defaults
│ │ ├── twam_base_cfg.py # released-checkpoint model config (20-d action, tactile)
│ │ ├── twam_posttrain_cfg.py # post-training recipe (edit this — POST_TRAINING.md)
│ │ ├── twam_posttrain_server_cfg.py # paired serve config (inherits the recipe)
│ │ └── twam_server_cfg.py # inference-server config (released pretrain ckpt)
│ ├── dataset/ # LeRobot latent datasets + bucket sampler
│ ├── distributed/ # FSDP2 sharding helpers
│ ├── utils/ # flow-matching scheduler, logging, websocket serving
│ │ └── Simple_Remote_Infer/ # websocket policy server + client
│ ├── train.py # training / post-training entry point
│ ├── n0_twam_server.py # inference server (obs → action)
│ └── render_mot.py # roll-out / visualization
├── script/ # data preparation for post-training
│ ├── encode_lerobot_n0_latents.py # RGB frames → Wan2.2 VAE latents
│ ├── encode_tactile_latent.py # tactile videos → global/local latents
│ ├── build_task_pool.py # task pool + norm stats (absEE / delta)
│ ├── make_serve_bundle.py # checkpoint + base model → serve bundle
│ └── build_segment_index.py # segment index (CSV) over multi-repo buckets
├── run_posttrain.sh # post-training launcher
├── diagrams/ · example_client/ # figures, example observations,
│ # simple_client.py (open-loop smoke),
│ # closed_loop_client.py (drive your robot)
├── docs/ # INSTALL.md · POST_TRAINING.md · DEPLOY.md
├── requirements.txt · pyproject.toml · LICENSE
See INSTALL.md. In short:
pip install .
pip install flash-attn --no-build-isolation| Model | Contents | Link |
|---|---|---|
|
|
transformer/ vae/ text_encoder/ tokenizer/ + norm_stat_pretrain.json + empty_emb.pt
|
NeoteAI/n0-twam-base |
| Post-trained · UniVTAC 8 tasks · absEE |
transformer/ + train_meta.json
|
NeoteAI/n0-twam-univtac-absee |
| Post-trained · UniVTAC 8 tasks · delta EE |
transformer/ + train_meta.json
|
NeoteAI/n0-twam-univtac-delta |
| Post-trained · NeoSim 12 tasks · absEE |
transformer/ + train_meta.json
|
NeoteAI/n0-twam-neosim-absee |
| Post-trained · NeoSim 12 tasks · delta EE |
transformer/ + train_meta.json
|
NeoteAI/n0-twam-neosim-delta |
The post-trained checkpoints are multi-task models, post-trained from
n0-twam-base with this toolkit's
recipe (POST_TRAINING.md): UniVTAC 8 = the 8
single-arm tasks inherited from UniVTAC, NeoSim 12 = the 12 NeoSim benchmark
tasks (4 single-arm + 8 dual-arm); absEE = absolute end-effector actions,
delta EE = horizon-delta actions; all trained on marker-less rgb tactile.
Each task trains with its own normalization stats, so serve them with the
multitask_server config, which selects the task and wires its stats, keys and
prompt — see DEPLOY.md.
from huggingface_hub import snapshot_download
bundle = snapshot_download("NeoteAI/n0-twam-base")import torch
from n0_twam.models.utils import load_mot_checkpoint
# `transformer/` holds the released checkpoint (config.json + weights).
model = load_mot_checkpoint(f"{bundle}/transformer",
torch_dtype=torch.bfloat16, torch_device="cuda")
print(f"{sum(p.numel() for p in model.parameters()) / 1e9:.2f} B params") # -> 7.16To adapt it to your own data, follow POST_TRAINING.md; to serve it, follow DEPLOY.md.
NeoSim is our vision–tactile simulation
benchmark. Its eval/ clients speak this server's streaming protocol natively:
the client runs the simulator, queries the server chunk by chunk, and re-grounds
the model's KV cache on the actually executed observations.
1. Set up NeoSim — follow the NeoSim installation guide. Keep it in its own environment: the simulator and the model server have incompatible dependency stacks, which is what the websocket split is for.
2. Launch the inference server (bundle + config: DEPLOY.md):
python -m n0_twam.n0_twam_server --config-name posttrain_server --port 296013. Run the evaluation client in the NeoSim repo (held-out protocol = 20 episodes, seeds 100–119):
python eval/eval_twam_ee_cl.py <task> demo --server_host <server-ip> --server_port 29601 \
--prompt "<training prompt, verbatim>" --start_seed 100 --total_num 20Key matching and tactile-representation alignment are covered in DEPLOY.md — Serving for NeoSim evaluation.
Not using NeoSim? example_client/closed_loop_client.py
speaks the same streaming protocol — reset, chunked inference, and KV-cache
re-grounding — without a simulator, so you can drive the policy from your own
robot. It needs only numpy, websockets and msgpack. See
DEPLOY.md — Close the loop.
- Three modality experts (MoT). Separate video, tactile, and action experts are coupled through shared cross-attention, each with its own width/FFN.
- One flow-matching objective for all modalities, with the three streams weighted equally (1 : 1 : 1) and a shared per-frame noise schedule, so video, tactile, and the action tokens of a frame stay temporally aligned and can be co-generated.
- Tactile as both a target and a condition. A global tactile stream is co-generated as a diffusion target (predicted as a residual over the first frame), while an optional local tactile pathway feeds the action expert the current observed tactile frame through cross-attention. During multi-task pretraining the tactile condition is randomly dropped (p=0.1) for robustness to missing sensors; single-task post-training keeps tactile always present.
@misc{n0twam2026,
title={$N_0$-TWAM: Scaling Tactile-Native World-Action Model for Contact-Rich Manipulation},
author={NeoteAI Team and Fudan TEAI Team},
year={2026},
eprint={2607.23783},
archivePrefix={arXiv},
url={https://arxiv.org/abs/2607.23783},
}This work builds upon several excellent open-source projects:
- LingBot-VA — the causal world-modeling framework this work builds on
- Wan2.2 — video diffusion transformer backbone (TI2V-5B) and the 48-channel VAE
- FastWAM — the Mixture-of-Transformers (shared-attention video + action experts) design our MoT references
- LeRobot — dataset format and tooling
Released under the CC-BY-NC-SA-4.0 license. See LICENSE. Redistributed third-party components keep their own licenses (the Wan2.2 VAE / text encoder / tokenizer and the openpi-derived websocket client are Apache-2.0 — their notices are retained).
