Friedrich-M/UniMatePublic

8 configs: 4 data combinations x 2 model variants, all at 60 frames.

The two axes are independent and all four combinations are implemented, so full x adaln and graph x cross_attn also run if you set them in a config; the two shipped pairings are the ones the paper compares.

Not shared across configs: training.batch_size, training.num_steps, model.num_layers and dataset.max_joints are tuned per data combination (GPU memory tracks batch x frames x joints; depth grows with the data: 6 layers for a single Truebones / Mixamo source, 8 for Objaverse, 10 for the full mixture). dataset.max_joints is 100 for truebones_* / mixamo_* and 60 for objaverse_* / uniml3d_*; dataset.min_joints is 5 everywhere. Together they bound the skeleton sizes a run admits (object types outside the range are dropped) and, through what survives, the joint-axis padding width. Mixamo is a single skeleton, so its configs also turn off the object-type balancing (a no-op with one type) and, by choice, the topology augmentations. Every other setting is identical.

Clips are drawn by a power-law-balanced sampler when training.balanced is set — a type with n clips is sampled in proportion to n^(1-sampler_alpha), so at the default sampler_alpha = 0.5 a species with 100 clips is seen ten times as often as one with a single clip rather than a hundred times — then augmented on the fly — joint addition, leaf removal, chain pooling and per-bone length perturbation (dataset.use_*_aug) — so the model sees more topologies than the data literally contains. Every clip is padded to max_joints on the joint axis and max_motion_length on the time axis, with masks carried alongside; nothing padded ever contributes to attention or to the loss.

Training is flow matching (training.diff_model = "flow"): the network predicts the velocity of a linear interpolant between noise and data, under a masked L2 loss plus two auxiliary terms computed on the reconstructed clean motion — a geodesic rotation loss (training.lambda_geo) and a velocity-smoothness loss (training.lambda_smooth). Conditioning is dropped with probability model.cond_mask_prob so the same weights serve the conditional and unconditional branches that classifier-free guidance interpolates at sampling time. AdamW with a cosine schedule and warmup, gradient clipping at training.max_grad_norm, and an EMA copy of the weights (training.use_ema) — the copy inference loads by default.

The text encoder (google/flan-t5-base by default) is fetched from the Hugging Face Hub on first use. Every run loads it once to embed all captions and joint names; pre-computing those embeddings beside the features keeps it out of the run entirely:

python -m unimate.tools.precompute_text_emb --config configs/uniml3d_60frames_graph_adaln.json

This writes caption_emb_cache.npz and joint_emb_cache.npz into each dataset/features/<dataset>/ the config uses. Captions are cached per token (the sequence cross_attn attends; adaln mean-pools it), joint names as one pooled vector each, keyed by the cleaned joint vocabulary that stage 3 produces — the shared naming is what lets the same anatomical joint embed identically across rigs. Re-run it after regenerating captions or joint names: anything the cache misses is still encoded at load time, so a stale cache costs speed rather than correctness.

A non-trivial share of the Objaverse-XL rigs and clips are defective: rest poses that lie flat, are rotated or are inverted, and clips that stitch several unrelated actions together. Training on them can destabilize or collapse a run, and isolated spikes in the training loss are usually a symptom of bad data rather than of optimization.

To localize the problem, first train on Mixamo and Truebones alone — set dataset.dataset_list to ["truebones", "mixamo"] in a copy of a config. If that run is healthy, the fault is on the Objaverse side. From there, inspect the skeleton preview videos of the suspect object types under dataset/features/objaverse/videos/, by eye or with an automated pass, and add the offending rigs and clips to the stage-4 skip lists that tools/patch_annotations.py maintains.

Test cases are a JSON map from <object_type>-<case_id> to a prompt. object_type must exist in the dataset; case_id is a free-form tag that names the output files:

{
  "Dog-walk": "a dog walks forward at a steady pace",
  "Dragon-takeoff": "a dragon flaps its wings and takes off"
}

--test_cases_json is itself optional: without it every clip of the dataset's eval split is enumerated as a test case (falling back to unique (object_type, caption) pairs on train when there is no eval split). --test_cases_txt (one object_type per line) drives unconditional sampling, which requires --cfg_scale 1.0.

Each run writes:

scripts/run_sample_motion_text.sh <exp_dir> [test_cases_json] [cfg_scale] runs the command above with the conda environment activated, the GPU with the most free memory selected, and an output directory named after the test-case file. Run it with -h for its options.

To animate the original mesh with a generated motion, hand the .npy files to stage 5 of the data pipeline, which exports an animated GLB + FBX:

bash scripts/run_animate_motion.sh objaverse \
    outputs/uniml3d_60frames_graph_adaln/samples/motions/Dog-walk-rep_0-0.npy \
    outputs/animated

It accepts several files or a directory, and reads the rig from dataset/features/<dataset>/cond.npy by default.

--keep_frames takes signed indices (negatives count back from the generation window), so "0,-1" fills in everything between a clip's first and last pose.

{ "mixamo-Squat-000": "A human squats and then rises back up" }

KEEP_FRAMES="0,-1" bash scripts/run_sample_motion_inbetween.sh \
    outputs/uniml3d_60frames_graph_adaln cases.json

--keep_joints matches case-insensitively against either the rig's own bone names or the cleaned vocabulary.

{ "<objaverse_uid>-turn-head-000": "The robot walks forward." }

KEEP_JOINTS="Hips,Spine,Neck,Head" bash scripts/run_sample_motion_edit.sh \
    outputs/uniml3d_60frames_graph_adaln cases.json

Test-case values become lists of prompts, one per segment; --expand_overlap sets how many frames consecutive segments share.

{ "mixamo-sequence": ["A human stands up.", "A human walks forward.", "A human turns around in place."] }

EXPAND_OVERLAP=10 bash scripts/run_sample_motion_expand.sh \
    outputs/uniml3d_60frames_graph_adaln cases.json

Each mode writes into its own subdirectory of --output_dir (inbetween/, motion_edit/, motion_expand/) alongside a small JSON recording the constraint that produced it, and each has a wrapper in scripts/ — run any of them with -h for the full option list.

In-betweening and editing clamp against a real clip, so their test-case keys must be <object_type>-<clip_id> naming a clip the dataset actually holds; that clip's motion is saved beside the result as <case_id>-gt_rep_<r>-<i>.npy for side-by-side comparison. --gt_start_frame pins which window of the clip is used instead of a random one. Editing trims both the sample and the GT to the clip's true length, while in-betweening generates the full window and trims only the GT — so align the two on frame 0 rather than assuming equal lengths. All three modes need --cfg_scale > 1.0 and are mutually exclusive with each other.