Control of chaotic systems has far-reaching implications in engineering, including fluid-based energy and transport systems, among many other fields. In real-world applications, control algorithms typically operate only with partial information about the system (partial observability) due to limited sensing, which leads to sub-optimal performance when compared to the case where a controller has access to the full system state (full observability). While it is well-known that the effect of partial observability can be mediated by introducing a memory component, which allows the controller to keep track of the system’s partial state history, the effect of the type of memory on performance in chaotic regimes is poorly understood. In this study we investigate the use of reinforcement learning for controlling chaotic flows using only partial observations. We use the chaotic Kuramoto–Sivashinsky equation with a forcing term as a model system. In contrast to previous studies, we consider the flow in a variety of dynamic regimes, ranging from mildly to strongly chaotic. We evaluate the loss of performance as the number of sensors available to the controller decreases. We then compare two different frameworks to incorporate memory into the controller, one based on recurrent neural networks and another novel mechanism based on transformers. We demonstrate that the attention-based framework robustly outperforms the alternatives in a range of dynamic regimes. In particular, our method yields improved control in highly chaotic environments, suggesting that attention-based mechanisms may be better suited to the control of chaotic systems.