jmportilla/ais-major-scratch-10k

jmportilla/ais-major-scratch-10k is an experimental AIS-style motion in-betweening checkpoint for the Animaj/MIB rig-controller dataset format. It was trained from scratch with a modified AIS-BiLSTM architecture designed to improve timing awareness and reduce noisy per-controller gating behavior.

This is not a generic Hugging Face Transformers model. It is a PyTorch checkpoint intended to be loaded through the motionib code in this project.

Model Purpose

The model predicts dense animation controller curves from sparse input keyposes. Given a sequence where only selected block/keypose frames are visible, the model fills the missing in-between frames for a 596-dimensional rig-controller vector.

The released AIS-BiLSTM baseline blends:

• an interpolation path between surrounding keyposes
• a synthesis path predicted from a Bi-LSTM hidden sequence
• a learned beta gate that blends interpolation and synthesis per controller dimension

This model keeps that core AIS idea but adds stronger timing and spatial gate structure.

Architectural Improvements Over The Base AIS-BiLSTM

Compared with AnimajSAS/AIS_BI_LSTM_v0, this checkpoint uses:

model.type = improved_ais_bilstm

Main changes:

• Explicit temporal progress conditioning Each frame receives four deterministic timing features: phase, segment_len, dist_prev, and dist_next. These tell the prediction heads where the frame sits between the previous and next observed keypose.

• Grouped beta gating The released AIS head predicts a separate beta gate for each of the 596 controller dimensions. This variant predicts 64 beta groups and expands them across the controller vector. The goal is to reduce noisy per-dimension gate switching and encourage related controller dimensions to choose interpolation-vs-synthesis behavior more coherently.

• AIS structure retained The model still predicts alpha, p_interp, p_synth, beta, and final pred, so it remains close to the AIS design while adding improved inductive bias.

Training objective changes:

• weighted L1 reconstruction
• velocity loss: 0.05
• acceleration loss: 0.01
• spectral loss: 0.001
• gate total-variation regularization: 0.001
• keypose reconstruction loss: 10.0

Architecture Configuration

model:
  type: improved_ais_bilstm
  pose_dim: 596
  input_dim: 597
  temporal_dim: 4
  hidden_size: 512
  num_layers: 2
  dropout: 0.3003
  synthesis_hidden: 1024
  beta_groups: 64

Training configuration:

max_steps: 10000
batch_size: 8
gradient_accumulation_steps: 8
optimizer: AdamW
lr: 0.0001
weight_decay: 0.001

The source checkpoint packaged here came from:

runs/ais_major_scratch_10k/checkpoints/best.pt

Full Benchmark Results

The table below comes from the local motionib.official_benchmark adapter using the full available public benchmark splits:

• held_out_algorithmic: 201 clips
• held_out_random: 201 clips, 90% random masking
• production: 56 clips

Lower is better for all listed metrics.

Important caveat: this adapter is designed to closely match the public protocol, but the upstream animaj-lab/mib-ais Lightning evaluator remains the final oracle for paper-level claims.

Test set	Model	Clips	Shifted distance	NPSS	Missing L1	Full L1
held_out_algorithmic	base	201	0.465652	2.182326	0.024197	0.021374
held_out_algorithmic	improved fine-tune	201	0.221540	0.897697	0.011201	0.009611
held_out_algorithmic	major scratch 10k	201	0.211539	0.901307	0.012029	0.010211
held_out_random	base	201	0.367792	1.359091	0.021008	0.018642
held_out_random	improved fine-tune	201	0.310235	1.326799	0.017078	0.015172
held_out_random	major scratch 10k	201	0.293350	1.266688	0.016796	0.014917
production	base	56	0.257322	0.671743	0.016206	0.012205
production	improved fine-tune	56	0.165956	0.485759	0.011167	0.008027
production	major scratch 10k	56	0.156249	0.493669	0.010951	0.007838

Improvement Over The Released Base Model

For the full available benchmark splits, major_scratch_10k improves over the released base checkpoint by:

Test set	Shifted distance	NPSS	Missing L1	Full L1
held_out_algorithmic	-54.6%	-58.7%	-50.3%	-52.2%
held_out_random	-20.2%	-6.8%	-20.1%	-20.0%
production	-39.3%	-26.5%	-32.4%	-35.8%

The scratch model is strongest on shifted distance, the timing-tolerant pose error metric used as the main paper-style comparison signal. On the full benchmark:

• major_scratch_10k has the best shifted distance on all three test sets.
• major_scratch_10k has the best NPSS on held_out_random.
• The fine-tuned improved model has slightly better NPSS on held_out_algorithmic and production.
• major_scratch_10k has the best missing-frame L1 on held_out_random and production.

Metric Definitions

• Shifted distance Paper-style shifted temporal L1. It compares each ground-truth segment to nearby predicted segments, reducing over-penalization from small timing shifts. Lower is better.

• NPSS Normalized Power Spectrum Similarity. It compares temporal frequency/rhythm distribution between prediction and ground truth. Lower is better.

• Missing L1 Mean absolute controller error on masked/in-between frames only. This isolates the frames the model actually has to predict.

• Full L1 Mean absolute controller error over the full sequence, including observed keypose frames.

Files

• model.pt PyTorch training checkpoint.

• training_config.yaml Full training configuration for this checkpoint.

• official_protocol_summary_metrics.csv Aggregate benchmark metrics for the full available benchmark splits.

• official_protocol_clip_metrics.csv Per-clip benchmark metrics.

• README.md This model card.

Loading

import torch
from motionib.models.improved_ais import build_improved_ais_bilstm

checkpoint = torch.load("model.pt", map_location="cpu", weights_only=False)
model = build_improved_ais_bilstm(checkpoint["config"]["model"], device="cpu")
model.load_state_dict(checkpoint["model"])
model.eval()

Expected Input Format

This model expects the tensor layout used by the motionib training code:

• input_seq: shape [batch, time, 597] 596 masked controller dimensions plus one missing-frame mask channel.

• prev_pose: shape [batch, time, 596] Previous observed keypose pose for each frame.

• next_pose: shape [batch, time, 596] Next observed keypose pose for each frame.

• temporal_features: shape [batch, time, 4] phase, segment_len, dist_prev, dist_next.

• observed_mask: optional shape [batch, time] Used to hard-copy observed keyposes at inference.

• keypose_values: optional shape [batch, time, 596] Ground-truth values at observed keypose frames.

Example forward pass:

with torch.no_grad():
    output = model(
        input_seq,
        prev_pose,
        next_pose,
        temporal_features=temporal_features,
        observed_mask=observed_mask,
        keypose_values=keypose_values,
    )

predicted_sequence = output["pred"]

Reproducing The Benchmark

From the project repository:

python -m motionib.official_benchmark \
  --config configs/train/ais_repro.yaml \
  --checkpoint base=artifacts/hf_model/model.safetensors \
  --checkpoint improved=runs/ais_improved_from_base_10k/checkpoints/best.pt \
  --checkpoint major_scratch=model.pt \
  --test-set all \
  --max-clips 0 \
  --device auto \
  --output-dir reports/official_protocol_compare_full

Limitations

• This is an experimental research checkpoint, not a production-ready animation system.
• It was trained for only 10,000 steps from scratch. Longer scratch training may improve results.
• The official upstream Lightning evaluator should still be used before making publication-level claims.
• The model is tied to the released Animaj/MIB controller-vector format and does not generalize to arbitrary rigs without retraining or adaptation.
• The benchmark compares controller-vector prediction quality, not rendered character quality or animator retake time.

Related Work

This model is based on the AIS motion in-betweening setup and the public AnimajSAS/AIS_BI_LSTM_v0 release. The architecture here is a local experimental variant that keeps the AIS interpolation/synthesis idea while adding temporal conditioning and grouped gate structure.