TacK-JEPA — pretrained checkpoints
Tactile Kinematic Joint-Embedding Predictive Architecture — a force-native, kinematically-grounded, action-conditioned world model for articulated multi-taxel tactile sensing (Allegro Hand). 100% simulation-sourced (Genesis physics); no physical hardware involved.
- Code / training pipeline: https://github.com/AmoghShrivastava/tack-jepa
- Design doc: see
PRD.mdin the code repo - License: MIT (weights + code). The training data used to produce these checkpoints is not released.
Input and output
Per simulation timestep:
- Input to the encoder: a graph of ~2,200 nodes, one per taxel (tactile sensing element) spread across the hand. Each node's features are exact 3D world-frame position (from forward kinematics — given, not learned or inferred), surface normal, a 3D force vector (normal + shear), and a link-ID embedding. Edges connect spatially-close taxels (recomputed every timestep) plus a static same-link backbone. The predictor is additionally conditioned on the hand's 22-dim action vector (16 finger joints + 6 floating-wrist pose).
- Output: a predicted latent vector k steps in the future — not an image, not raw sensor values. An online encoder embeds the current graph; the predictor takes that embedding plus the action sequence and predicts what a target encoder (EMA copy of the online encoder, gradients stopped) will produce k steps later. Training minimizes the distance between the two; the model never predicts or reconstructs a raw taxel value directly.
- Small frozen probe heads (evaluation only, not part of the training loss) map the encoder's latent to interpretable quantities — per-taxel force magnitude, a binary slip indicator, contact area — to check whether the latent space captures useful physical structure.
How this differs from standard tactile models
JEPA-style pretraining for touch already exists (Meta's SPARSH trains and releases an I-JEPA variant), and action-conditioned tactile world models already exist too (TacForeSight, Dream-Tac) — so neither "JEPA for touch" nor "world model for touch" alone is the novel claim here. The gap this project targets is more specific:
| Sensor / signal | Geometry | Prediction target | |
|---|---|---|---|
| SPARSH, AnyTouch, T3, UniT | Optical tactile image (GelSight/DIGIT), 1–2 patches/hand | Inferred from pixels | Latent prediction within a static image |
| TacForeSight, Dream-Tac | Real force/torque sensor(s), 1–2 points | Not kinematically tracked | Future latent state, action-conditioned |
| TacK-JEPA | Simulated distributed taxel field, ~2,200 sensors/hand | Exact, via forward kinematics | Future latent state, action-conditioned |
What's in this repo
Five checkpoints from the full training sweep — the §7.2 ablation set, each trained for an equal 6000-step budget on the full dataset (16 object variants, press/grasp/slide trajectories, object-disjoint train/val split):
| Directory | Variant |
|---|---|
baseline_stagec/ | Full model: graph encoder + forward-kinematics taxel positions + VICReg |
no_fk_stagec/ | Ablation: forward-kinematics taxel-position grounding removed |
image_native_stagec/ | Ablation: image-based (ViT) encoder instead of the graph encoder |
reconstruction_stagec/ | Ablation: reconstruction objective instead of JEPA prediction |
no_vicreg_stagec/ | Ablation: VICReg regularization removed |
Each directory contains checkpoint.pt, metrics.jsonl (training curve), and
resolved_config.yaml (the exact config used for that run).
Usage
Load with the model classes in the code repo (models/encoder.py,
models/predictor.py, models/ablations/) — checkpoints are plain PyTorch
state_dict bundles, not tied to any particular serving framework.
hf download AmoghShrivastava1/tack-jepa --local-dir ./tack-jepa-checkpoints
Caveats
This is an independent research project, not a production system. See the
code repo's ROADMAP.md for the full decisions log, known limitations, and
open investigations (e.g. an observed AUROC gap between the graph-encoder and
image-native variants, tracked but not fully closed as of the last update).