tack-jepa

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.md in 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:

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):

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).