MAGE: Multi-scale Autoregressive Generation for Offline Reinforcement Learning

Chenxing Lin; Xinhui Gao; Haipeng Zhang; Xinran Li; Haitao Wang; Songzhu Mei; Chenglu Wen; Weiquan Liu

cs.LGFeb 27, 2026

MAGE: Multi-scale Autoregressive Generation for Offline Reinforcement Learning

Authors: Chenxing Lin, Xinhui Gao, Haipeng Zhang, Xinran Li, Haitao Wang, Songzhu Mei, Chenglu Wen, Weiquan Liu,

Organizations: Fujian Key Laboratory of Urban Intelligent Sensing and Computing, School of Informatics, Xiamen University (XMU), China · Key Laboratory of Multimedia Trusted Perception and Efficient Computing, XMU, China · Hong Kong University of Science and Technology · School of Computer, National University of Defense Technology, China ·

Abstract

Generative models have gained significant traction in offline reinforcement learning (RL) due to their ability to model complex trajectory distributions. However, existing generation-based approaches still struggle with long-horizon tasks characterized by sparse rewards. Some hierarchical generation methods have been developed to mitigate this issue by decomposing the original problem into shorter-horizon subproblems using one policy and generating detailed actions with another. While effective, these methods often overlook the multi-scale temporal structure inherent in trajectories, resulting in suboptimal performance. To overcome these limitations, we propose MAGE, a Multi-scale Autoregressive GEneration-based offline RL method. MAGE incorporates a condition-guided multi-scale autoencoder to learn hierarchical trajectory representations, along with a multi-scale transformer that autoregressively generates trajectory representations from coarse to fine temporal scales. MAGE effectively captures temporal dependencies of trajectories at multiple resolutions. Additionally, a condition-guided decoder is employed to exert precise control over short-term behaviors. Extensive experiments on five offline RL benchmarks against fifteen baseline algorithms show that MAGE successfully integrates multi-scale trajectory modeling with conditional guidance, generating coherent and controllable trajectories in long-horizon sparse-reward settings.

Figures & tables

Figure 1: Schematic Illustration of Generation-based Offline RL Methods. (a) Decision Transformer follows a step-by-step, autoregressive generation process. (b) Decision Diffuser utilizes an all-at-once, denoising-based generation approach. (c) MAGE operates in a top-down manner, first establishing a macroscopic outline of a trajectory and then progressively refining it with microscopic details.

Figure 2: MAGE Overview: (a) Multi-scale Representation: hierarchical quantization of trajectories across scales for global–local structure modeling. (b) Condition-guided Decision-making: autoregressive latent prediction with conditional refinement for consistent trajectory generation.

Figure 3: A Maze Game: The dark grid cells are walls. The trajectories are plotted in red. Lighter color represents earlier timesteps. The red circles shown in (c) are subgoals.

Scenario		IQL	DT	ADT	CARP	DD	D-QL	HDMI	HD	MAGE
Pen	Expert	128.0±9.2	116.3±1.2	113.3±12.1	112.7±19.8	107.6±7.6	112.6±8.1	109.5±8.0	121.4±14.3	147.8±4.9
	Human	78.4±8.2	67.6±5.4	70.1±16.1	62.3±21.4	64.1±9.0	66.0±8.3	66.2±8.8	47.6±14.9	137.1±9.0
	Cloned	83.4±8.1	64.4±1.4	35.9±13.1	12.5±15.2	47.7±9.2	49.3±8.0	48.3±8.9	13.9±9.7	108.4±17.6
Door	Expert	106.6±0.3	104.8±0.3	105.1±0.1	98.4±4.7	87.0±0.8	93.7±0.8	85.9±0.9	105.9±0.6	106.8±0.1
	Human	3.2±1.8	4.4±0.8	7.5±2.3	5.0±4.6	6.9±1.2	8.0±1.2	7.1±1.1	0.2±0.0	16.5±0.9
	Cloned	3.0±1.7	7.6±3.2	1.8±1.3	0.0±0.0	9.0±1.6	10.6±1.7	9.3±1.6	4.5±3.6	20.5±2.5

Table 1: The Average Scores for the Adroit Scenarios.

Scenario		IQL	DT	ADT	CARP	DD	DQL	HDMI	HD	MAGE
Kitchen	Partial	59.7±8.3	31.4±19.5	64.2±5.1	32.5±2.6	65.0±2.8	60.5±6.9	-	73.3±1.4	91.3±3.2
Kitchen	Mixed	53.2±1.6	25.8±5.0	69.2±3.3	30.0±2.2	57.0±2.5	62.6±5.1	69.2±1.8	71.7±2.5	86.3±3.3
Average		56.5	28.6	66.7	31.3	61.0	61.6	-	72.5	88.8

Table 2: The Average Scores for the Franka Kitchen Scenarios.

Scenario		BC	CQL	IQL	DT	ADT	DD	D-QL	HD	MAGE
Diverse	U-maze	47.2±4.0	37.2±3.7	70.6±3.7	51.7±0.4	83.0±3.1	49.2±3.1	66.2±8.6	94.0±4.9	95.2±2.2
	Medium	0.8±0.8	67.2±3.5	61.7±6.1	0.0±0.0	83.4±1.9	4.0±2.8	78.6±10.3	88.7±8.1	98.2±1.3
	Large	0.0±0.0	20.5±13.2	27.6±7.8	0.0±0.0	65.4±4.9	0.0±0.0	56.6±7.6	83.6±5.8	84.6±3.6
Play	U-maze	55.2±4.1	92.7±1.9	83.3±4.5	57.0±9.8	83.8±2.3	73.1±2.5	93.4±3.4	72.2±2.0	92.2±2.7
	Medium	0.0±0.0	65.7±11.6	64.6±4.9	0.0±0.0	82.0±1.7	8.0±4.3	76.6±10.8	42.0±1.9	92.0±2.7
	Large	0.0±0.0	20.7±7.2	42.5±6.5	0.0±0.0	71.0±1.3	0.0±0.0	46.4±8.3	54.7±2.0	75.8±4.3

Table 3: The Average Scores for the Antmaze Scenarios.

Scenario		CQL	IQL	DT	ADT	CARP	DD	HDMI	HD	MAGE
Maze2D	U-maze	-8.9±6.1	42.1±0.5	31.0±21.3	60.5±2.0	26.2±3.9	116.2±2.7	120.1±2.5	128.4±3.6	145.4±3.2
	Medium	86.1±9.6	34.8±2.7	8.2±4.4	109.4±6.2	65.8±2.4	122.3±2.1	121.8±1.6	135.6±3.0	155.0±3.3
	Large	23.7±36.7	61.7±3.5	2.3±0.9	155.4±10.4	0.7±2.0	125.9±1.6	128.6±2.9	155.8±2.5	159.4±2.9
Single-task Average		33.6	46.2	13.8	108.4	30.9	121.5	123.5	139.9	153.3
Multi2D	U-maze	25.4±5.8	13.5±3.0	15.6±2.4	66.9±5.2	82.5±3.4	128.2±2.1	131.3±1.8	144.1±1.2	150.4±1.8
	Medium	8.3±3.9	8.3±3.4	6.3±1.6	108.5±6.2	48.7±2.6	129.7±2.7	131.6±1.9	140.2±1.6	147.7±3.1

Table 4: The Average Scores for the Maze2D and Multi2D Scenarios.

Scenario	ADT	HD	1	2	4	6	8	10
Pen-Expert	113.3±12.1	121.4±14.3	123.5±9.1	127.5±5.2	134.2±7.7	139.5±5.7	147.8±4.9	149.9±9.2
Door-Cloned	1.8±1.3	4.5±3.6	5.2±1.8	6.0±2.1	10.7±2.3	14.0±2.7	20.5±2.5	17.0±2.7

Table 5: Performance of varying the number of temporal scales

K

Scenario	ADT	HD	S	A	A+CQL	R, S, A	Ours
Pen-Expert	113.3±12.1	121.4±14.3	127.7±3.3	127.1±13.0	127.6±4.6	124.9±7.8	147.8±4.9
Door-Cloned	1.8±1.3	4.5±3.6	11.9±2.7	9.1±2.3	4.9±1.0	17.2±3.0	20.5±2.5

Table 6: Different trajectory sequence generation schemes.

Scenario	ADT	HD	$\cancel{R}$ in $\mathcal{D}$	$\cancel{R}$ in $m_{k>1}$	$\cancel{R}$ in $\mathcal{L}_{cond}$	Ours
Pen-Expert	113.3 ± 12.1	121.4±14.3	140.3±9.1	139.5±6.7	139.9±7.2	147.8±4.9
Door-Cloned	1.8 ± 1.3	4.5±3.6	12.3±1.9	16.3±2.6	17.1±2.5	20.5±2.5

Table 7: Ablation results of condition and constraint loss on the Adroit scenarios

Method	Ours	DT	TT	ADT	DD	HD
Time (ms)	27.30±0.69	6.49±0.11	12863.07±19.37	7.81±0.14	2339.16±11.37	1480.21±25.18

Table 8: Average inference time (ms) in Adroit environments.

Appendix figures & tables20 assets

Supplementary material from the paper’s appendix.

Appendix

No.	Algorithm	Brief Description
1	BC ( Bain and Sammut, 1995 )	Trains a model by directly learning from examples provided by an expert, enabling the model to mimic the expert’s behavior
2	CQL 1 1 1 https://github.com/BY571/CQL ( Kumar et al., 2020 )	Updates Q-values conservatively to improve stability and sample efficiency
3	IQL 2 2 2 https://github.com/ikostrikov/implicit_q_learning ( Kostrikov et al., 2022 )	Decouples policy updates from Q-value estimation to improve the stability and performance of offline reinforcement learning
4	DT 3 3 3 https://github.com/kzl/decision-transformer ( Chen et al., 2021 )	Uses a Transformer architecture to model sequences of states, actions, and rewards, enabling it to make decisions based on the entire history of interactions and desired outcomes
5	TAP 4 4 4 https://github.com/ZhengyaoJiang/latentplan ( Zhang et al., 2023 )	Addresses the challenge of high-dimensional control by planning over temporally abstract latent actions, drastically reducing decision latency while improving performance
6	ADT 5 5 5 https://github.com/mamengyiyi/Autotuned-Decision-Transformer ( Ma et al., 2024 )	Jointly optimizes high-level prompt and low-level action policies to improve Decision Transformer’s ability to stitch trajectories

Appendix

Table 9: Baseline algorithms

Hyper-parameter	Value
learning rate	2e-4
horizon	24
dropout rate	0.1
discount	0.99
codebook size	512
transformer blocks	8

Appendix

Table 10: Hyperparameter Settings for Adroit environment

Hyper-parameter	Value
learning rate	2e-4
horizon	24
dropout rate	0.1
discount	0.99
codebook size	2048
transformer blocks	8

Appendix

Table 11: Hyperparameter Settings for Franka Kitchen environment

Hyper-parameter	Value
learning rate	2e-4
horizon	24
dropout rate	0.1
discount	0.998
codebook size	1024
transformer blocks	8

Appendix

Table 12: Hyperparameter Settings for Antmaze environment

Hyper-parameter	Value
learning rate	2e-4
horizon	24
dropout rate	0.1
discount	0.99
codebook size	256
transformer blocks	4

Appendix

Table 13: Hyperparameter Settings for Maze2D and Multi2D environment

Hyper-parameter	Value
learning rate	2e-4
horizon	24
dropout rate	0.1
discount	0.99
codebook size	512
transformer blocks	8

Appendix

Table 14: Hyperparameter Settings for Gym locomotion control tasks

Scenario		BC	CQL	TT	TAP	MAGE
Pen	Expert	94.6±3.2	-1.4±2.3	101.8±13.8	127.4±7.7	147.8±4.9
	Human	71.0±6.2	13.7±16.9	2.0±3.4	76.5±8.5	137.1±9.0
	Cloned	51.9±15.1	1.0±6.6	38.8±13.3	57.4±8.7	108.4±17.6
Door	Expert	105.1±2.4	-0.3±0.0	101.6±4.8	104.8±0.8	106.8±0.1
	Human	2.6±5.7	5.5±1.3	0.1±0.0	8.8±1.1	16.5±0.9
	Cloned	-0.1±0.0	-0.3±0.0	0.0±0.0	11.7±1.5	20.5±2.5

Appendix

Table 15: Additional baseline comparisons for the Adroit scenarios. Results are averaged over 5 random training seeds, with each seed tested 20 times. Bold numbers indicate the best performance.

Scenario		BC	CQL	Diffuser	MAGE
Kitchen	Partial	41.3±3.7	51.3±7.7	52.5±2.5	91.3±3.2
Kitchen	Mixed	48.9±0.7	51.3±7.7	55.7±1.3	86.3±3.3
Average		45.1	51.3	54.1	88.8

Appendix

Table 16: Additional baseline comparisons for the Franka Kitchen Scenarios. Results are averaged over 5 random training seeds, with each seed tested 20 times. Bold numbers show the best performance.

Scenario		MPPI	Diffuser	RGG	MAGE
Maze2D	U-maze	33.2	113.9±3.1	108.8±1.4	145.4±3.2
	Medium	10.2	121.5±2.7	131.8±0.5	155.0±3.3
	Large	5.1	123.0±6.4	135.4±1.7	159.4±2.9
Single-task Average		16.2	119.5	125.3	153.3
Multi2D	U-maze	41.2	128.9±1.8	128.3±0.8	150.4±1.8
	Medium	15.4	127.2±3.4	130.0±0.9	147.7±3.1

Appendix

Table 17: Additional baseline comparisons for the Maze2D and Multi2D scenarios. Results are averaged over 5 random training seeds, with each seed tested 100 times. Bold numbers indicate the best performance.

Task	Dataset	BC	CQL	IQL	DT	TT	CARP	TAP	Diffuser	HD	MAGE
HalfCheetah	Medium-Expert	55.2	91.6	86.7	86.8	95.0 $\pm$ 0.2	57.1 $\pm$ 2.7	91.8 $\pm$ 0.8	88.9 $\pm$ 0.3	92.5 $\pm$ 0.3	95.2 $\pm$ 0.2
	Medium	42.6	49.2	47.4	42.6	46.9 $\pm$ 0.4	38.2 $\pm$ 1.4	45.0 $\pm$ 0.1	42.8 $\pm$ 0.3	46.7 $\pm$ 0.2	43.9 $\pm$ 0.2
	Md-Replay	36.6	45.5	44.2	36.6	41.9 $\pm$ 2.5	34.6 $\pm$ 0.6	40.8 $\pm$ 0.6	37.7 $\pm$ 0.5	38.1 $\pm$ 0.7	46.0 $\pm$ 0.2
Walker2d	Medium-Expert	107.5	108.8	109.6	108.1	101.9 $\pm$ 6.8	102.2 $\pm$ 2.8	107.4 $\pm$ 0.9	106.9 $\pm$ 0.2	107.1 $\pm$ 0.1	110.3 $\pm$ 0.1
	Medium	75.3	83.0	78.3	74.0	79.0 $\pm$ 2.8	60.7 $\pm$ 1.8	64.9 $\pm$ 2.1	79.6 $\pm$ 0.6	84.0 $\pm$ 0.6	83.5 $\pm$ 0.5
	Md-Replay	32.3	77.2	73.9	79.4	82.6 $\pm$ 6.9	42.7 $\pm$ 1.1	66.8 $\pm$ 3.1	70.6 $\pm$ 1.6	84.1 $\pm$ 2.2	87.8 $\pm$ 2.0

Appendix

Table 18: Baseline comparisons across different Gym locomotion control tasks. Bold numbers indicate the best performance.

Figure 4: Effect of Conditional Constraints on MAGE: without $\mathcal{L}_{\text{cond}}$ , the agent can still locate coins and reach the goal, but the agent generates the wrong trajectory that crosses the wall.

Scenario	Explicit	Latent (Ours)
Pen-Expert	136.5±6.7	147.8±4.9
Door-Cloned	0.4±0.3	20.5±2.5

Appendix

Table 19: Comparison between explicit and latent inverse dynamics models.

Scenario	128	256	512	1024
Pen-Expert	112.9±19.0	146.5±3.1	147.8±4.9	145.3±1.9
Door-Cloned	14.4±2.9	19.7±3.4	20.5±2.5	11.3±2.1

Appendix

Table 20: Effect of codebook size on performance.

Scenario	2	4	8	16
Pen-Expert	109.8±18.3	124.1±17.1	147.8±4.9	145.6±3.4
Door-Cloned	12.4±2.2	14.8±2.4	20.5±2.5	18.4±2.5

Appendix

Table 21: Effect of Transformer depth on performance.

Scenario	MAGE (direct decoder)	MAGE
Pen-Expert	132.4 $\pm$ 8.3	147.8 $\pm$ 4.9
Door-Cloned	12.0 $\pm$ 1.7	20.5 $\pm$ 2.5

Appendix

Table 22: Effect of the adapter module for applying

\mathcal{L}_{cond}

Scenario	MAGE (w/o condition)	MAGE (reweight)	MAGE
Pen-Expert	$92.1\pm 13.2$	$131.7\pm 10.4$	$147.8\pm 4.9$
Door-Cloned	$0.0\pm 0.0$	$6.1\pm 1.8$	$20.5\pm 2.5$

Appendix

Table 23: Effect of RTG-based reweighting compared with RTG conditioning.

Scenario	Diffuser	ADT	HIQL	MAGE
pointmaze-giant-navigate-v0	$0\pm 0$	$19\pm 4$	$46\pm 9$	$\mathbf{52\pm 5}$
antmaze-giant-navigate-v0	$2\pm 1$	$26\pm 4$	$\mathbf{65\pm 5}$	$58\pm 5$
antmaze-teleport-navigate-v0	$8\pm 3$	$23\pm 4$	$42\pm 3$	$\mathbf{49\pm 5}$
humanoidmaze-giant-navigate-v0	$0\pm 0$	$2\pm 1$	$12\pm 4$	$\mathbf{17\pm 4}$

Appendix

Table 24: Performance on OGBench long-horizon maze tasks.

Scenario	ADT	HD	MAGE (fixed config)	MAGE
Maze2d-Medium	109.4±6.2	135.6±3.0	146.9±4.2	155.0±3.3
Multi2d-Medium	108.5±6.2	140.2±1.6	142.1±1.8	147.7±3.1
Antmaze-Medium-Play	82.0±1.7	42.0±1.9	90.0±3.0	92.0±2.7
Antmaze-Medium-Diverse	83.4±1.9	88.7±8.1	94.8±2.2	98.2±1.3
Kitchen-Mixed	69.2±3.3	71.7±2.5	80.0±5.8	86.3±3.3
Kitchen-Partial	64.2±5.1	73.3±1.4	86.3±4.1	91.3±3.2

Appendix

Table 25: Performance with the default MAGE configuration.

Scenario	DT	K=8 w/o $L_{cond}$	K=1 w/o $L_{cond}$	K=1	K=2	K=4	K=8
Pen-Expert	$116.3\pm 1.2$	$136.6\pm 7.4$	$119.5\pm 11.5$	$123.5\pm 9.1$	$127.5\pm 5.2$	$134.2\pm 7.7$	$147.8\pm 4.9$
Door-Cloned	$7.6\pm 3.2$	$16.6\pm 2.1$	$3.4\pm 2.6$	$5.2\pm 1.8$	$6.0\pm 2.1$	$10.7\pm 2.3$	$20.5\pm 2.5$
Hammer-Expert	$117.4\pm 6.6$	$128.3\pm 0.3$	$113.2\pm 6.2$	$116.5\pm 0.6$	$121.7\pm 0.4$	$127.9\pm 0.3$	$131.7\pm 0.2$
Relocate-Expert	$104.2\pm 0.4$	$108.9\pm 1.7$	$100.1\pm 2.9$	$101.3\pm 3.8$	$102.5\pm 1.8$	$105.9\pm 1.4$	$109.6\pm 1.6$

Appendix

Table 26: Effect of temporal scales and the condition-guided decoder module.

	ADT	HDMI	HD	CARP	MAGE
Level/scales	Two-level	Two-level	Two-level	Multi-scale	Multi-scale
Number of Policies	2	2	2	1	1
Generation Model	Transformer	Diffusion	Diffusion	Transformer	Transformer
Generated Data at the first level	G	G	G	latent of A	latent of (S,R)
Generated Data at the last level	S,A	S,A	S,A	latent of A	latent of (S,R)
Condition at the first level	S	S,R	S,R	S	S,R

Appendix

Table 27: Comparison of different methods.

Offline RL

Offline reinforcement learning learns decision-making policies from fixed datasets without interacting with the environment during training.

See all 299 papers

Chain-of-Goals Hierarchical Policy for Long-Horizon Offline Goal-Conditioned RL
Feb 3, 2026Jinwoo Choi, Sang-Hyun Lee, Seung-Woo SeoReinforcement Learning Hierarchical RL
Department of Electrical and Computer Engineering, Seoul National University, Seoul, South Korea · Department of Automotive Engineering, Ajou University, Gyeonggi-do, South Korea
Offline RL with Hierarchical Action Chunking
Jul 23, 2026Ahad JawaidHierarchical RL Offline RL
The University of Texas at Dallas
HorizonFlow: Variable-Length Planning for Offline Goal-Conditioned RL
Sep 29, 2026JunHyeok Oh, Zian Jang, Byung-Jun LeeHierarchical Planning Trajectory Generation
Korea University. · Gauss Labs Inc., Seoul, Republic of Korea.

Feb 3, 2026J/K move · Enter open · S save

Chain-of-Goals Hierarchical Policy for Long-Horizon Offline Goal-Conditioned RL

Jinwoo Choi, Sang-Hyun Lee, Seung-Woo Seo

Department of Electrical and Computer Engineering, Seoul National University, Seoul, South Korea · Department of Automotive Engineering, Ajou University, Gyeonggi-do, South Korea

Reinforcement Learning Hierarchical RL Offline RL Long-Horizon Agent Tasks Goal-Conditioned RL Long-Horizon RL

Open paper AskPDF arXiv

Offline goal-conditioned reinforcement learning remains challenging for long-horizon tasks. While hierarchical approaches mitigate this issue by decomposing tasks, most existing methods rely on separate high- and low-level networks and generate only a single intermediate subgoal, leaving several structural limitations in long-horizon decision-making. To address this limitation, we draw inspiration from chain-of-thought reasoning and propose the Chain-of-Goals Hierarchical Policy (CoGHP), a novel framework that reformulates hierarchical decision-making as autoregressive sequence modeling within a unified architecture. Given a state and a final goal, CoGHP autoregressively generates a sequence of latent subgoals followed by the primitive action, where each latent subgoal acts as a reasoning step that conditions subsequent predictions. To implement this efficiently, we introduce an MLP-Mixer backbone, which supports cross-token communication and captures structural relationships among state, goal, latent subgoals, and action. Across challenging navigation and manipulation benchmarks, CoGHP consistently outperforms strong offline baselines, demonstrating improved performance on long-horizon tasks. Project page: https://wlsdn9350.github.io/projects/coghp/

Abstract

Offline RL

Explore similar work

Chain-of-Goals Hierarchical Policy for Long-Horizon Offline Goal-Conditioned RL

Offline RL with Hierarchical Action Chunking

HorizonFlow: Variable-Length Planning for Offline Goal-Conditioned RL