Research map / Reinforcement learning / Offline RL

Unveiling Markov heads in Pretrained Language Models for Offline Reinforcement Learning

ICML 2025 · Reinforcement learning › Offline RL · approach: Decision Transformer sequence models · 0 citations (OpenAlex)

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Identifies "Markov heads" in GPT-2 attention that focus on the last input token, proves they persist under random embeddings and fine-tuning, and proposes GPT2-DTMA (mixture of attention over heads) to adapt PLM-initialized decision transformers to both short- and long-term offline RL tasks.

Problem

Pretrained language models transferred to offline RL via decision transformers perform well in short-term (MuJoCo locomotion) environments but underperform a DT trained from scratch in long-term (PointMaze) environments; the knowledge responsible for this transfer is unclear.

Approach

Defines Markov heads as attention heads whose W^q (W^k)^T is a Markov matrix (dominant positive diagonal), shows theoretically that such heads keep last-token attention under random embeddings and small-step fine-tuning, and adds a Mixture-of-Attention gate over heads so the model can down-weight Markov heads in long-term tasks.

Key results

• GPT2-DTMA reaches 79.9 on Walker2d-medium and 77.4 on Hopper-medium, comparable to or above DT (74.3 and 67.4), in short-term MuJoCo tasks.
• On PointMaze-large, GPT2-DTMA reduces episode length to 203.0 from 257.3 for GPT2-DT, though it remains behind DT (195.3).
• Penalizing Markov head weights (GPT2-P) lowers short-term performance, supporting the claim that Markov heads carry the transferred benefit.

Measured effects reported in the paper

Surprising findings

• GPT-J, despite the same next-token prediction loss as GPT-2, does not exhibit the Markov matrix structure, and CLIP-DT and GPTJ-DT do not match GPT2-DT performance.
• Plain DTMA (MoA on a randomly initialized DT) gives no gain over DT, suggesting MoA helps only when Markov heads from PLM are present.

Limitations

• GPT2-DTMA does not outperform DT in long-term PointMaze environments, only narrows the gap.
• Markov head identification relies on a hand-chosen threshold r=8 and a limited set of environments; theory is stated with probabilistic bounds rather than tight guarantees.

Experimental setup

Datasets: D4RL MuJoCo Locomotion datasets D4RL PointMaze datasets

Benchmarks: D4RL MuJoCo Locomotion D4RL PointMaze

Metrics: normalized score episode length

Baselines: DT CQL Decision ConvFormer (DC) GPT2-DT (LaMo) DTMA CLIP-DT GPTJ-DT GPT2-P

Builds on: Decision Transformer LaMo Mixture of Experts / mixture of attention CQL

Closest papers by meaning

• Unleashing the Power of Pre-trained Language Models for Offline Reinforcement Learning (ICLR 2024 · similarity 0.90)
• Efficient Reinforcement Learning with Large Language Model Priors (ICLR 2025 · similarity 0.89)
• Zero-shot Model-based Reinforcement Learning using Large Language Models (ICLR 2025 · similarity 0.88)
• ShiQ: Bringing back Bellman to LLMs (NeurIPS 2025 · similarity 0.88)
• Q-SFT: Q-Learning for Language Models via Supervised Fine-Tuning (ICLR 2025 · similarity 0.87)
• On Predictability of Reinforcement Learning Dynamics for Large Language Models (ICLR 2026 · similarity 0.87)