README.md

# ReMoMask: Retrieval-Augmented Masked Motion Generation<br>

This is the official repository for the paper:
> **ReMoMask: Retrieval-Augmented Masked Motion Generation**
>
> Zhengdao Li\*, Siheng Wang\*, [Zeyu Zhang](https://steve-zeyu-zhang.github.io/)\*<sup>†</sup>, and [Hao Tang](https://ha0tang.github.io/)<sup>#</sup>
>
> \*Equal contribution. <sup>†</sup>Project lead. <sup>#</sup>Corresponding author.
>
> ### [Paper](https://arxiv.org/abs/2508.02605) | [Website](https://aigeeksgroup.github.io/ReMoMask) | [Model](https://huggingface.co/lycnight/ReMoMask) | [HF Paper](https://huggingface.co/papers/2508.02605)


# ✏️ Citation

```
@article{li2025remomask,
  title={ReMoMask: Retrieval-Augmented Masked Motion Generation},
  author={Li, Zhengdao and Wang, Siheng and Zhang, Zeyu and Tang, Hao},
  journal={arXiv preprint arXiv:2508.02605},
  year={2025}
}
```

---

# 👋 Introduction
Retrieval-Augmented Text-to-Motion (RAG-T2M) models have demonstrated superior performance over conventional T2M approaches, particularly in handling uncommon and complex textual descriptions by leveraging external motion knowledge. Despite these gains, existing RAG-T2M models remain limited by two closely related factors: (1) coarse-grained text-motion retrieval that overlooks the hierarchical structure of human motion, and (2) underexplored mechanisms for effectively fusing retrieved information into the generative process. In this work, we present **ReMoMask**, a structure-aware RAG framework for text-to-motion generation that addresses these limitations. To improve retrieval, we propose **Hierarchical Bidirectional Momentum** (HBM) Contrastive Learning, which employs dual contrastive objectives to jointly align global motion semantics and fine-grained part-level motion features with text. To address the fusion gap, we first conduct a systematic study on motion representations and information fusion strategies in RAG-T2M, revealing that a 2D motion representation combined with cross-attention-based fusion yields superior performance. Based on these findings, we design **Semantic Spatial-Temporal Attention** (SSTA), a motion-tailored fusion module that more effectively integrates retrieved motion knowledge into the generative backbone. Extensive experiments on HumanML3D, KIT-ML, and SnapMoGen demonstrate that ReMoMask consistently outperforms prior methods on both text-motion retrieval and text-to-motion generation benchmarks.



## TODO List

- [x] Upload our paper to arXiv and build project pages.
- [x] Upload the code.
- [x] Release TMR model.
- [x] Release T2M model.

# 🤗 Prerequisite
<details> 
<summary>details</summary>
  
## Environment
```bash
conda create -n remomask python=3.10
pip install torch==2.1.0 torchvision==0.16.0 torchaudio==2.1.0 --index-url https://download.pytorch.org/whl/cu118
pip install -r requirements.txt
conda activate remomask
```
We tested our environment on both A800 and H20.

## Dependencies
### 1. pretrained models
Dwonload the models from [HuggingFace](https://huggingface.co/lycnight/ReMoMask), and place them like:

```
remomask_models.zip
    ├── checkpoints/              #   Evaluation Models and Gloves
    ├── Part_TMR/
    │   └── checkpoints/        # RAG pretrained checkpoints
    ├── logs/                   # T2M pretrained checkpoints
    ├── database/               # RAG database
    └── ViT-B-32.pt            # CLIP model
```

### 2. Prepare training dataset 
Follow the instruction in [HumanML3D](https://github.com/EricGuo5513/HumanML3D.git), then place the result dataset to `./dataset/HumanML3D`.
</details>

# 🚀 Demo
<details> 
<summary>details</summary>
  
```bash
python demo.py \
    --gpu_id 0 \
    --ext exp_demo \
    --text_prompt "A person is playing the drum set." \
    --checkpoints_dir logs \
    --dataset_name humanml3d \
    --mtrans_name pretrain_mtrans \
    --rtrans_name pretrain_rtrans
# change pretrain_mtrans and pretrain_rtrans to your mtrans and rtrans after your training done
```
explanation:
* `--repeat_times`: number of replications for generation, default `1`.
* `--motion_length`: specify the number of poses for generation.

output will be in `./outputs/`
</details> 


# 🛠️ Train your own models
<details>
<summary>details</summary>
  
## Stage1: train a Motion Retriever
```bash
python Part_TMR/scripts/train.py \
    device=cuda:0 \
    train=train \
    dataset.train_split_filename=train.txt \
    exp_name=exp \
    train.optimizer.motion_lr=1.0e-05 \
    train.optimizer.text_lr=1.0e-05 \
    train.optimizer.head_lr=1.0e-05
# change the exp_name to your rag name
```
then build a rag database for training t2m model:
```bash
python build_rag_database.py \
    --config-name=config \
    device=cuda:0 \
    train=train \
    dataset.train_split_filename=train.txt \
    exp_name=exp_for_mtrans
```
you will get `./database`


##  Stage2: train a Retrieval Augmented Mask Model

### tarin a 2D RVQ-VAE Quantizer
```bash
bash run_rvq.sh \
    vq \
    0 \
    humanml3d \
    --batch_size 256 \
    --num_quantizers 6 \
    --max_epoch 50 \
    --quantize_dropout_prob 0.2 \
    --gamma 0.1 \
    --code_dim2d 1024 \
    --nb_code2d 256
# vq means the save dir
# 0 means gpu_0
# humanml3d means dataset
# change the vq_name to your vq name
```

### train a 2D Retrieval-Augmented Mask Transformer
```bash
bash run_mtrans.sh \
    mtrans \
    1 \
    0 \
    11247 \
    humanml3d \
    --vq_name pretrain_vq \
    --batch_size 64 \
    --max_epoch 2000 \
    --attnj \
    --attnt \
    --latent_dim 512 \
    --n_heads 8 \
    --train_split train.txt \
    --val_split val.txt
# 1 means using one gpu
# 0 means using gpu_0
# 11247 means ddp master port
# change the mtrans to your mtrans name
```


### train a 2D Retrieval-Augmented Residual Transformer
```bash
bash run_rtrans.sh \
    rtrans \
    2 \
    humanml3d \
    --batch_size 64 \
    --vq_name pretrain_vq \
    --cond_drop_prob 0.01 \
    --share_weight \
    --max_epoch 2000 \
    --attnj \
    --attnt
# here, 2 means cuda:0,1
# --vq_name: the vq model you want to use
# change the rtrans to your vq rtrans
```

</details>



# 💪 Evalution
<details>
<summary>details</summary>
  
## Evaluate the RAG  
```bash
python Part_TMR/scripts/test.py \
    device=cuda:0 \
    train=train \
    exp_name=exp_pretrain
# change exp_pretrain to your rag model
```


## Evaluate the T2M

### 1. Evaluate the 2D RVQ-VAE Quantizer
```bash
python eval_vq.py \
--gpu_id 0 \
--name pretrain_vq \
--dataset_name humanml3d \
--ext eval \
--which_epoch net_best_fid.tar
# change pretrain_vq to your vq
```

### 2. Evaluate the 2D Retrieval-Augmented Masked Transformer
```bash
python eval_mask.py \
    --dataset_name humanml3d \
    --mtrans_name pretrain_mtrans \
    --gpu_id 0 \
    --cond_scale 4 \
    --time_steps 10 \
    --ext eval \
    --repeat_times 1 \
    --which_epoch net_best_fid.tar
# change pretrain_mtrans to your mtrans
```


### 3. Evaluate the 2D Residual Transformer
HumanML3D:
```bash
python eval_res.py \
    --gpu_id 0 \
    --dataset_name humanml3d \
    --mtrans_name pretrain_mtrans \
    --rtrans_name pretrain_rtrans \
    --cond_scale 4 \
    --time_steps 10 \
    --ext eval \
    --which_ckpt net_best_fid.tar \
    --which_epoch fid \
    --traverse_res
# change pretrain_mtrans and pretrain_rtrans to your mtrans and rtrans
```
</details>



# 🤖 Visualization
<details>
<summary>details</summary>
  
## 1. download and set up blender
<details>
<summary>details</summary>
You can download the blender from [instructions](https://www.blender.org/download/lts/2-93/). Please install exactly this version. For our paper, we use `blender-2.93.18-linux-x64`. 
> 
### a. unzip it:
```bash
tar -xvf blender-2.93.18-linux-x64.tar.xz
```

### b. check if you have installed the blender successfully or not:
```bash
cd blender-2.93.18-linux-x64
./blender --background --version
```
you should see: `Blender 2.93.18 (hash cb886axxxx built 2023-05-22 23:33:27)`
```bash
./blender --background --python-expr "import sys; import os; print('\nThe version of python is ' + sys.version.split(' ')[0])"
```
you should see: `The version of python is 3.9.2`

### c. get the blender-python path
```bash
./blender --background --python-expr "import sys; import os; print('\nThe path to the installation of python is\n' + sys.executable)"
```
you should see: `	The path to the installation of python is /xxx/blender-2.93.18-linux-x64/2.93/python/bin/python3.9s`

### d. install pip for blender-python
```bash
/xxx/blender-2.93.18-linux-x64/2.93/python/bin/python3.9 -m ensurepip --upgrade
/xxx/blender-2.93.18-linux-x64/2.93/python/bin/python3.9 -m pip install --upgrade pip
```

### e. prepare env for blender-python
```bash 
/xxx/blender-2.93.18-linux-x64/2.93/python/bin/python3.9 -m pip install numpy==2.0.2
/xxx/blender-2.93.18-linux-x64/2.93/python/bin/python3.9 -m pip install matplotlib==3.9.4
/xxx/blender-2.93.18-linux-x64/2.93/python/bin/python3.9 -m pip install hydra-core==1.3.2
/xxx/blender-2.93.18-linux-x64/2.93/python/bin/python3.9 -m pip install hydra_colorlog==1.2.0
/xxx/blender-2.93.18-linux-x64/2.93/python/bin/python3.9 -m pip install moviepy==1.0.3
/xxx/blender-2.93.18-linux-x64/2.93/python/bin/python3.9 -m pip install shortuuid==1.0.13
/xxx/blender-2.93.18-linux-x64/2.93/python/bin/python3.9 -m pip install natsort==8.4.0
/xxx/blender-2.93.18-linux-x64/2.93/python/bin/python3.9 -m pip install pytest-shutil==1.8.1
/xxx/blender-2.93.18-linux-x64/2.93/python/bin/python3.9 -m pip install tqdm==4.67.1
/xxx/blender-2.93.18-linux-x64/2.93/python/bin/python3.9 -m pip install tqdm==1.17.0
```
</details>


## 2. calulate SMPL mesh:
```bash
python -m fit --dir new_test_npy --save_folder new_temp_npy --cuda cuda:0
```

## 3. render to video or sequence
```bash
/xxx/blender-2.93.18-linux-x64/blender --background --python render.py -- --cfg=./configs/render_mld.yaml --dir=test_npy --mode=video --joint_type=HumanML3D
```
- `--mode=video`: render to mp4 video
- `--mode=sequence`: render to a png image, calle sequence.

</details>

# 👍 Acknowlegements
We sincerely thank the open-sourcing of these works where our code is based on:

[MoMask](https://github.com/EricGuo5513/momask-codes),
[MoGenTS](https://github.com/weihaosky/mogents),
[ReMoDiffuse](https://github.com/mingyuan-zhang/ReMoDiffuse),
[MDM](https://github.com/GuyTevet/motion-diffusion-model),
[TMR](https://github.com/Mathux/TMR),
[ReMoGPT](https://ojs.aaai.org/index.php/AAAI/article/view/33044)

## 🔒 License
This code is distributed under an [CC BY-NC-SA 4.0](https://creativecommons.org/licenses/by-nc-sa/4.0/deed.en).

Note that our code depends on other libraries, including CLIP, SMPL, SMPL-X, PyTorch3D, and uses datasets that each have their own respective licenses that must also be followed.