Views
No views yet
1# Run under repository: https://github.com/thu-ml/RDT2
2
3import torch
4import numpy as np
5from models.normalizer import LinearNormalizer
6from vqvae.models.multivqvae import MultiVQVAE
7
8# Load from the Hub (replace with your repo id once published)
9vae = MultiVQVAE.from_pretrained("outputs/vqvae_hf").cuda().eval()
10normalizer = LinearNormalizer.load(
11 "<Path_to_normalizer>" # Download from:
12 # http://ml.cs.tsinghua.edu.cn/~lingxuan/rdt2/umi_normalizer_wo_downsample_indentity_rot.pt
13)
14
15# Load your RELATIVE action chunk
16action_chunk = torch.zeros((1, 24, 20))
17# action_chunk shape: (B, T, action_dim)
18# - T = 24: predicts the future 0.8s in fps=30 → 24 frames
19# - action_dim = 20: following UMI setting (both arms, right to left)
20# - [0-2]: RIGHT ARM end effector position in (x, y, z), unit: m
21# - [3-8]: RIGHT ARM end effector rotation (6D representation)
22# - [9]: RIGHT ARM gripper width, normalized to [0, 0.088], 0.088 means fully open
23# - [10-12]: LEFT ARM end effector position in (x, y, z), unit: m
24# - [13-18]: LEFT ARM end effector rotation (6D representation)
25# - [19]: LEFT ARM gripper width, normalized to [0, 0.088], 0.088 means fully open
26
27# Normalize action
28nsample = normalizer["action"].normalize(action_chunk).cuda()
29
30# Encode → tokens
31# tokens: torch.LongTensor with shape (B, num_valid_action_token)
32# num_valid_action_token = 27, values in range [0, 1024)
33tokens = vae.encode(nsample) # or vae.encode(action_chunk)
34
35# Decode back to continuous actions
36recon_nsample = vae.decode(tokens)
37recon_action_chunk = normalizer["action"].unnormalize(recon_nsample)
381@article{liu2026rdt2,
2 title={RDT2: Exploring the Scaling Limit of UMI Data Towards Zero-Shot Cross-Embodiment Generalization},
3 author={Liu, Songming and Li, Bangguo and Ma, Kai and Wu, Lingxuan and Tan, Hengkai and Ouyang, Xiao and Su, Hang and Zhu, Jun},
4 journal={arXiv preprint arXiv:2602.03310},
5 year={2026}
6}