CroSTAta introduces a novel State Transition Attention (STA) mechanism designed to improve robotic manipulation tasks by better capturing temporal dependencies and state transitions in sequential data. It enables policies to adapt behavior based on execution history, particularly in failure and recovery patterns.

The Cross-State Transition Attention Transformer modulates standard attention weights based on learned state evolution patterns. This approach is combined with temporal masking during training to encourage temporal reasoning from historical context, achieving more than 2x improvement over standard cross-attention on precision-critical tasks.

To train the model on a specific environment (e.g., PegInsertionSide-v1) using ManiSkill:

To evaluate a saved checkpoint:

Note: sl_agent inference uses by default the policy's method predict_batch (batch prediction). For efficient inference, use the policy's method predict (prediction with cache).

The model was evaluated using specific recovery datasets available on the Hugging Face Hub: