abyssrz/dino-wm-libero

DINO-WM Safety Filter for LIBERO

A learned safety filter for robot manipulation based on DINOv2 and a world model. Trained on LIBERO-Object tasks with backdoor trigger detection.

Files

File	Size	Description
best_classifier.pth	258 MB	World model backbone + safety classifier head
policy.pth	18 MB	HJ safety critic (DDPG policy trained in WM latent space)
libero_action_stats.npz	554 B	Action normalization statistics

BD3 checkpoints

The BD3 artifacts are available in the bd3/ directory. They are stored separately so the previously published root-level checkpoints remain unchanged. See bd3/README.md for file descriptions and SHA-256 checksums.

Architecture

• Vision encoder: DINOv2 ViT-S/14 (registers variant, 384-dim features, 224×224 input)
• World model: VideoTransformer (6 layers, 16 heads, MLP dim 2048)
  • Context window: 3 frames → predict 1 frame
  • Action embedding: 10-dim (from 7-DoF delta actions)
  • State input: 8-dim proprioception (end-effector + gripper)
  • Predicts: front camera features, wrist camera features, next state, failure score
• Safety classifier: MLP head on transformer features → scalar risk score per frame
• HJ critic: DDPG policy scoring (state, action) pairs in WM latent space

Usage

Load the model via dino_safety_filter.py from the latent-safety-dino repository:

from dino_safety_filter import build_libero_dino_safety_filter

safety_filter = build_libero_dino_safety_filter(
    wm_ckpt_path="best_classifier.pth",
    hj_policy_path="policy.pth",
    action_stats_path="libero_action_stats.npz",
    threshold=0.0,
    device="cuda:0",
)
# safety_filter.filter_action(action) -> (filtered_action, accepted)

Training Data

• Dataset: LIBERO-Object (10 tasks, ~500 demonstrations)
• Labels: Per-timestep failure annotations (0 = safe, 1 = unsafe/backdoor)
• Backdoor trigger: Malicious scene objects + predicate violations

Training Pipeline

1. train_dino_decoder.py — VQVAE decoder for visual reconstruction
2. train_dino_wm.py — World model dynamics (front/wrist/state prediction heads)
3. train_dino_classifier.py — Safety classifier (failure_head only, backbone frozen)
4. train_critic.sh — HJ avoid critic in WM latent space (DDPG)

Citation

Based on DINO-WM: https://arxiv.org/abs/2411.04983