| --- |
| license: apache-2.0 |
| library_name: pytorch |
| tags: |
| - robotics |
| - robot-manipulation |
| - self-supervised |
| - masked-autoencoder |
| - egocentric-video |
| - visual-representation |
| - pytorch |
| pipeline_tag: feature-extraction |
| --- |
| |
| # STP: Spatiotemporal Predictive Pre-training for Robotic Motor Control |
|
|
| Pre-trained visual encoder for robotic manipulation, from the paper published in |
| *International Journal of Computer Vision* (2026): |
| [10.1007/s11263-026-02948-3](https://link.springer.com/article/10.1007/s11263-026-02948-3) |
|
|
| Official code: [MCG-NJU/STP](https://github.com/MCG-NJU/STP) |
|
|
| STP pre-trains a plain image ViT on large-scale egocentric video (Ego4D / |
| EgoClip), with no language, action, depth or teacher supervision. It jointly |
| learns content and motion features by reconstructing the masked current frame |
| and the masked future frame, using a shared encoder with decoupled dual |
| decoders. |
|
|
| At evaluation the encoder is frozen, both decoders are discarded, and the |
| `[CLS]` token is used as the visual representation for downstream policy |
| learning. |
|
|
| ## Files |
|
|
| Both checkpoints are the ViT-B/16 model pre-trained on the **EgoClip** subset |
| of Ego4D (3.8 M clips) for 50 epochs. |
|
|
| | file | size | contents | |
| |---|---|---| |
| | `stp_vitb.pth` | ~558 MiB | model weights (encoder + both decoders), 146.34 M parameters | |
| | `stp_vitb_encoder.pth` | ~327 MiB | decoder-free encoder, for use with a plain timm ViT | |
| | `config.json` | small | architecture and pre-training recipe | |
| | `models_stp.py`, `util/` | small | model definition, so this repo is self-contained | |
| | `loader.py` | small | minimal `STPEncoder` wrapper | |
|
|
| ## Architecture |
|
|
| | | | |
| |---|---| |
| | encoder | ViT-B/16, 12 layers, 12 heads, 768 dim, fixed sincos positional embedding | |
| | spatial decoder | 8 blocks, 16 heads, 512 dim, transformer encoder blocks | |
| | temporal decoder | 8 blocks, 16 heads, 512 dim, self-attention + cross-attention + FFN | |
| | input | 224 × 224, ImageNet normalisation | |
|
|
| The spatial branch masks the current frame at 75% and reconstructs it. The |
| temporal branch masks the future frame at 95% and reconstructs it, conditioned |
| on the current-frame feature `Z_c` used as a fixed key/value, so the future-frame |
| query cannot write back into the current-frame representation. |
|
|
| ## Usage |
|
|
| ```python |
| import torch |
| from models_stp import mae_vit_base_patch16 |
| from huggingface_hub import hf_hub_download |
| |
| path = hf_hub_download('yangjiange/STP', 'stp_vitb.pth') |
| model = mae_vit_base_patch16() |
| model.load_state_dict(torch.load(path, map_location='cpu')['model'], strict=True) |
| model.eval() |
| |
| # images: (N, 3, 224, 224), ImageNet-normalised -> (N, 768) |
| feats = model.forward_features(images) |
| ``` |
|
|
| Concatenate multi-view and multi-frame features in the channel dimension |
| together with the robot's proprioceptive state, then train the policy. STP does |
| not constrain the policy architecture. |
|
|
| ## Citation |
|
|
| ```bibtex |
| @Article{Yang2026STP, |
| author = {Jiange Yang and Bei Liu and Jianlong Fu and Bocheng Pan |
| and Gangshan Wu and Limin Wang}, |
| journal = {International Journal of Computer Vision}, |
| title = {Spatiotemporal Predictive Pre-training for Robotic Motor Control}, |
| year = {2026}, |
| volume = {134}, |
| doi = {10.1007/s11263-026-02948-3}, |
| } |
| ``` |
|
|
| ## License |
|
|
| Apache-2.0. |
|
|
