KV-Control

Sparse-keyframe, multi-joint controllable text-to-motion generation.

This repository releases the reference implementation of the paper main method — a KV-Control adapter on top of the T-Concat v4 masked-transformer backbone — for the HumanML3D text-to-motion benchmark. The base masked transformer learns text-to-motion generation; the KV-Control adapter adds a 1.5 M-parameter low-rank Key/Value projection on top of the frozen backbone that lets you steer the generated motion through sparse 3-D keyframe targets on any subset of the 22 SMPL joints (single joint or multi-joint, arbitrary keyframe spacing).

text prompt: "a person walks forward while drawing a heart with both hands"
+ sparse 3-D keyframes: pelvis trajectory + 2 wrist heart-curve targets
→ KV-Control → 22-joint sequence → SMPL mesh

Repository layout

KV_control/
├── kvctrl/                       # python package
│   ├── models/                   # base + KV-Control + VQ + eval encoder
│   │   ├── mask_transformer/     # T-Concat v4 + KV-Control adapter
│   │   └── vq/                   # part-aware VQ encoder
│   ├── data/                     # HumanML3D dataset wrappers
│   ├── utils/                    # motion ops, metrics, plot helpers
│   ├── motion_loaders/           # eval-time data loaders
│   ├── visualize/                # SMPLify joints→mesh
│   ├── visualization/            # BVH / animation helpers
│   ├── generation/               # avoidance loss, control sampling
│   ├── options/                  # argparse option groups
│   ├── common/                   # skeleton & quaternion helpers
│   └── paths.py                  # central repo-relative path registry
├── scripts/                      # CLI entry points
│   ├── train_base.py             # train T-Concat v4 base model
│   ├── train_kvctrl.py           # fine-tune KV-Control adapter on base
│   ├── eval_kvctrl.py            # multi-joint sparse-control evaluation
│   ├── eval_base_cfg.py          # text-only CFG eval of the base model
│   ├── demo_designed_trajectories.py  # 8 designed-trajectory demos
│   ├── sanity_check_equivalence.py    # end-to-end inference smoke test
│   ├── fit_smpl.py / prep_mesh_pkl.py / render_video.py
│   └── render_blender/render.py  # Blender storyboard / video renderer
├── checkpoints/                  # populated by download_checkpoints.sh
│   ├── base_t_concat_v4/model/net_best_fid.tar               # 168 MB  frozen base
│   ├── kv_control/model/net_best_top3.tar                    # 520 MB  cross multi-joint (\KPSmulti=0.80)
│   ├── kv_control_trajectory/model/net_best_kps.tar          # 520 MB  single-joint pelvis headline (\KPSours=0.40)
│   ├── vqvae/{net_best_fid.pth, skeleton_partition.json}
│   ├── stats/{mean.npy, std.npy}
│   ├── clip/ViT-B-32.pt
│   └── t2m/{Comp_v6_KLD005, text_mot_match, length_estimator}
├── dataset/HumanML3D             # (symlink) — see "Dataset" below
├── body_models/smpl/             # SMPL neutral mesh + face / J_regressor
├── glove/                        # GloVe-style vocab for length estimator
└── docs/, assets/

All file paths used by the code live in kvctrl/paths.py; every script resolves them relative to the repository root. Override individual locations with environment variables KVCTRL_CKPT_ROOT, KVCTRL_DATASET_ROOT, KVCTRL_BODY_MODELS_ROOT, MASKCONTROL_CLIP_MODEL_PATH.

Install

git clone https://github.com/CHDTevior/KV-Control.git
cd KV-Control

# Conda environment (Python 3.10, CUDA 11.8+)
conda create -n kvctrl python=3.10
conda activate kvctrl
pip install -r requirements.txt
pip install -e .

Dataset

This project uses HumanML3D (Guo et al., 2022). Follow the upstream HumanML3D recipe to produce the standard split:

HumanML3D/
├── new_joint_vecs/         # 263-dim feature vectors per sequence (.npy)
├── new_joints/             # 22-joint positions per sequence (.npy)
├── texts/                  # one .txt per sequence
├── train.txt, val.txt, test.txt
├── Mean.npy, Std.npy

Once you have a local HumanML3D directory, expose it to the project:

ln -sfn /absolute/path/to/HumanML3D dataset/HumanML3D
# (or set KVCTRL_DATASET_ROOT in the environment)

We do not redistribute the dataset. The link is mandatory because all training and eval scripts read from dataset/HumanML3D.

Checkpoints

All released checkpoints — base T-Concat v4, KV-Control adapter, VQ-VAE, normalization stats, the frozen text-motion eval encoder, and the CLIP ViT-B/32 visual backbone — live on Hugging Face at Tevior/KV-Control. The download script writes them into the canonical sub-tree under checkpoints/:

bash scripts/download_checkpoints.sh

huggingface-cli reads HF_TOKEN if present; the released repo is public so a token is not required.

Quickstart

1. End-to-end equivalence smoke test

python scripts/sanity_check_equivalence.py \
    --out output/equivalence/new_joints.npy

Generates 120 frames for the walk-heart-both-hands designed trajectory using the released cross KV-Control checkpoint, saves the joints, and prints the keyframe error (KPS) on that single hand-designed 6-joint, 8-keyframe trajectory. This is a generation smoke test, not a benchmark: the KPS here is ≈ 1.7 cm (expected for one hard designed sample — it is not the paper metric, and the script does not itself diff against an external reference). For the paper test-set numbers (\KPSours, \KPSmulti) use the HumanML3D evaluation in §3.

2. Designed-trajectory demos (8 patterns)

python scripts/demo_designed_trajectories.py \
    --traj walk_heart_both_hands \
    --motion_length 196 \
    --keyframe_stride 8 \
    --control_set 4joint \
    --out output/demos/walk_heart_both_hands/joints.npy

Supported --traj values: walks_arc_arms_high, walk_S_arms_alternate, walk_spiral_arms_swing, walk_heart_both_hands, heart_both_hands_stationary, circle_arms_high, walk_wave, walk_zigzag_arms_alternate.

3. Reproduce the paper metrics on the HumanML3D test split

Both paper headline checkpoints are released. The M3 hybrid protocol below is the exact paper evaluation protocol (Stage-1 dynamic TTT each_iter=35 --ttt_dynamic, T=10; Stage-2 600-step embedding optimisation, cfg=3.25, --cond_drop_prob 0.0, --pred_num_batch 16 --seed 3407, HumanML3D test).

Single-joint pelvis — paper Tab 4 headline row (\KPSours):

python scripts/eval_kvctrl.py \
    --ckpt checkpoints/kv_control_trajectory/model/net_best_kps.tar \
    --control trajectory --split test --repeat_times 5 \
    --time_steps 10 --cond_scale 3.25 \
    --each_iter 35 --each_lr 6e-2 --ttt_dynamic \
    --last_iter 600 --last_lr 6e-2 --cond_drop_prob 0.0 \
    --pred_num_batch 16 --seed 3407 --nb_code 128 \
    --vq_checkpoint checkpoints/vqvae/net_best_fid.pth \
    --vq_partition_file checkpoints/vqvae/skeleton_partition.json \
    --trans_path checkpoints/base_t_concat_v4/model/net_best_fid.tar
# Expected (5r mean): KPS ≈ 0.40 cm, FID ≈ 0.065, Top-3 ≈ 0.799

Multi-joint cross — paper Tab 4 multi-joint block (\KPSmulti):

python scripts/eval_kvctrl.py \
    --ckpt checkpoints/kv_control/model/net_best_top3.tar \
    --control cross --split test --repeat_times 5 \
    --time_steps 10 --cond_scale 3.25 \
    --each_iter 35 --each_lr 6e-2 --ttt_dynamic \
    --last_iter 600 --last_lr 6e-2 --cond_drop_prob 0.0 \
    --pred_num_batch 16 --seed 3407 --nb_code 128 \
    --vq_checkpoint checkpoints/vqvae/net_best_fid.pth \
    --vq_partition_file checkpoints/vqvae/skeleton_partition.json \
    --trans_path checkpoints/base_t_concat_v4/model/net_best_fid.tar
# Expected (5r mean): KPS ≈ 0.80 cm (best 0.71)

A single --repeat_times 1 run lands within the 5-repeat 95% CI of the values above; use 5 repeats for paper-grade numbers.

4. Train base T-Concat v4 from scratch

torchrun --nproc_per_node=4 scripts/train_base.py \
    --dataset_name t2m \
    --name my_base_t_concat_v4 \
    --transformer_variant t_concat_v4 \
    --batch_size 64 --max_epoch 2000 \
    --lr 2e-4 --warm_up_iter 2000

5. Fine-tune KV-Control adapter on a pre-trained base

torchrun --nproc_per_node=4 scripts/train_kvctrl.py \
    --dataset_name t2m \
    --name my_kv_control \
    --trans_path checkpoints/base_t_concat_v4/model/net_best_fid.tar \
    --control cross --batch_size 256 \
    --kv_rank 64 --max_epoch 500 \
    --xent 0.5 --ctrl_loss 0.5

VQ-VAE provenance

The part-aware VQ-VAE tokenizer (checkpoints/vqvae/net_best_fid.pth) was trained in a separate research project ("part-aware-vqvae") that is not part of this release. We ship the inference-side VQ code under kvctrl/models/{vqvae,encdec,quantize_cnn,resnet}.py that is bit-exactly compatible with the released checkpoint — see docs/VQ_PROVENANCE.md for architecture / hyper-parameter details and how to re-train the VQ if needed.

Method (one paragraph)

The base masked transformer (20 layers × 8 heads, d=384, ff=1536) tokenises HumanML3D motions with a part-aware VQ-VAE (128 codes × 6 parts) and learns text-conditional codebook prediction with sparse cross-attention to a CLIP text adapter. The KV-Control adapter freezes the entire base; for each cross-attention block it adds two low-rank projections (rank 64) that compute a sparse 3-D Key/Value pair from any subset of the 22 joints' target positions. At inference time we sample with classifier-free guidance and run test-time refinement on the predicted joints to tighten control-point tracking.

Equivalence to the development repo

The released inference code was validated to reproduce the development repository's KV-Control outputs under a fixed seed, checkpoint and prompt. scripts/sanity_check_equivalence.py regenerates the designed walk-heart-both-hands trajectory and reports its keyframe error so an install can be sanity-checked end-to-end (it generates joints and prints KPS; it does not itself diff against an external reference).

License

MIT. See LICENSE.

Citation

Please update CITATION.bib once the paper is accepted.

@article{kvcontrol2026,
  title  = {KV-Control: Sparse-Keyframe Multi-Joint Text-to-Motion Generation},
  author = {... (anonymous during review) ...},
  year   = {2026},
  note   = {Under review at SIGGRAPH Asia.}
}

Acknowledgements

This code builds on MaskControl, HumanML3D, CLIP (OpenAI ViT-B/32), the SMPLify body model, StableMoFusion's CLIP-text-adapter recipe, and the MaskControl/T2M-GPT training infrastructure.