# CroSTAta: Cross-State Transition Attention Transformer for Robotic Manipulation

arXiv: https://arxiv.org/abs/2510.00726 · HTML: https://mlsanity.com/paper/2510.00726/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in Robot manipulation imitation learning, one of 13 directions in Offline RL and imitation learning from data.
- Builds on 33 of its 38 references that are in the map.
- First posted on arXiv 1 Oct 2025.

## Placed in

- Offline RL and imitation learning from data / [Robot manipulation imitation learning](https://mlsanity.com/offlinerl/robot-manipulation-imitation-learning/)

## Builds on (most cited first, 30 of 33)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Diffusion Policy: Visuomotor Policy Learning via Action Diffusion](https://mlsanity.com/paper/2303.04137/) (2023)
- [π0: A Vision-Language-Action Flow Model for General Robot Control](https://mlsanity.com/paper/2410.24164/) (2024)
- [Learning Fine-Grained Bimanual Manipulation with Low-Cost Hardware](https://mlsanity.com/paper/2304.13705/) (2023)
- [Decision Transformer: Reinforcement Learning via Sequence Modeling](https://mlsanity.com/paper/2106.01345/) (2021)
- [GR00T N1: An Open Foundation Model for Generalist Humanoid Robots](https://mlsanity.com/paper/2503.14734/) (2025)
- [What Matters in Learning from Offline Human Demonstrations for Robot Manipulation](https://mlsanity.com/paper/2108.03298/) (2021)
- [RLBench: The Robot Learning Benchmark & Learning Environment](https://mlsanity.com/paper/1909.12271/) (2019)
- [Offline Reinforcement Learning as One Big Sequence Modeling Problem](https://mlsanity.com/paper/2106.02039/) (2021)
- [Temporal Convolutional Networks for Action Segmentation and Detection](https://mlsanity.com/paper/1611.05267/) (2016)
- [VIMA: General Robot Manipulation with Multimodal Prompts](https://mlsanity.com/paper/2210.03094/) (2022)
- [RoboCasa: Large-Scale Simulation of Everyday Tasks for Generalist Robots](https://mlsanity.com/paper/2406.02523/) (2024)
- [ManiSkill2: A Unified Benchmark for Generalizable Manipulation Skills](https://mlsanity.com/paper/2302.04659/) (2023)
- [MimicGen: A Data Generation System for Scalable Robot Learning using Human Demonstrations](https://mlsanity.com/paper/2310.17596/) (2023)
- [Causal Confusion in Imitation Learning](https://mlsanity.com/paper/1905.11979/) (2019)
- [ManiSkill: Generalizable Manipulation Skill Benchmark with Large-Scale Demonstrations](https://mlsanity.com/paper/2107.14483/) (2021)
- [Behavior Generation with Latent Actions](https://mlsanity.com/paper/2403.03181/) (2024)
- [A Survey of Imitation Learning: Algorithms, Recent Developments, and Challenges](https://mlsanity.com/paper/2309.02473/) (2023)
- [Data Scaling Laws in Imitation Learning for Robotic Manipulation](https://mlsanity.com/paper/2410.18647/) (2024)
- [Self-Supervised Visual Planning with Temporal Skip Connections](https://mlsanity.com/paper/1710.05268/) (2017)
- [What Matters in Language Conditioned Robotic Imitation Learning over Unstructured Data](https://mlsanity.com/paper/2204.06252/) (2022)
- [HAMSTER: Hierarchical Action Models For Open-World Robot Manipulation](https://mlsanity.com/paper/2502.05485/) (2025)
- [Fighting Copycat Agents in Behavioral Cloning from Observation Histories](https://mlsanity.com/paper/2010.14876/) (2020)
- [In-Context Imitation Learning via Next-Token Prediction](https://mlsanity.com/paper/2408.15980/) (2024)
- [Feedback in Imitation Learning: The Three Regimes of Covariate Shift](https://mlsanity.com/paper/2102.02872/) (2021)
- [Needle In A Multimodal Haystack](https://mlsanity.com/paper/2406.07230/) (2024)
- [Bidirectional Decoding: Improving Action Chunking via Guided Test-Time Sampling](https://mlsanity.com/paper/2408.17355/) (2024)
- [SkillMimicGen: Automated Demonstration Generation for Efficient Skill Learning and Deployment](https://mlsanity.com/paper/2410.18907/) (2024)
- [Reflective Planning: Vision-Language Models for Multi-Stage Long-Horizon Robotic Manipulation](https://mlsanity.com/paper/2502.16707/) (2025)
- [Autonomous Improvement of Instruction Following Skills via Foundation Models](https://mlsanity.com/paper/2407.20635/) (2024)

## Abstract (arXiv)

Learning robotic manipulation policies through supervised learning from demonstrations remains challenging when policies encounter execution variations not explicitly covered during training. While incorporating historical context through attention mechanisms can improve robustness, standard approaches process all past states in a sequence without explicitly modeling the temporal structure that demonstrations may include, such as failure and recovery patterns. We propose a Cross-State Transition Attention Transformer that employs a novel State Transition Attention (STA) mechanism to modulate standard attention weights based on learned state evolution patterns, enabling policies to better adapt their behavior based on execution history. Our approach combines this structured attention with temporal masking during training, where visual information is randomly removed from recent timesteps to encourage temporal reasoning from historical context. Evaluation in simulation shows that STA consistently outperforms standard attention approach and temporal modeling methods like TCN and LSTM networks, achieving more than 2x improvement over cross-attention on precision-critical tasks. The source code and data can be accessed at https://github.com/iit-DLSLab/croSTAta