Gaze-Regularized Vision Transformers for Imitation Learning

Behavioral cloning policies trained on state-action pairs alone are prone to causal confusion: in visually rich environments they latch onto features that correlate with the expert's action during training but do not cause it, and fail at test time. GABRIL (Banayeeanzade et al., 2025) addresses this by regularizing a convolutional policy's activation map toward human gaze masks. We study whether the same idea transfers to Vision Transformers, whose attention maps give a more direct handle on where the model looks. We train a ViT-Small behavioral cloning policy on Atari demonstrations with human eye-tracking data, and add a loss that aligns the last layer's CLS-to-patch attention with the expert's gaze heatmap. We compare two design choices in the regularizer, the loss used to match attention to gaze (distribution-level cross-entropy vs. per-pixel binary cross-entropy) and whether attention heads are aggregated before or after upsampling to frame resolution, against an unregularized baseline. On Seaquest, all four gaze-regularized variants score 180, compared with 80 for the unregularized ViT, while the choice of matching loss and aggregation order makes no difference to the score.