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

AI summary

Overview

Research area: Offline goal-conditioned reinforcement learning (offline GCRL), specifically hierarchical policy learning for long-horizon tasks, with borrowings from chain-of-thought reasoning and the MLP-Mixer sequence architecture.

Technical level: Advanced. The paper assumes familiarity with offline RL (value functions, advantage-weighted regression, implicit Q-learning), hierarchical decomposition into high-level subgoals and low-level control, and sequence-model architectures.

Scope in one sentence: The paper proposes CoGHP, a unified autoregressive model that generates a chain of latent subgoals followed by a primitive action, and evaluates it on OGBench navigation and manipulation benchmarks against eight offline baselines.

Authors are Jinwoo Choi and Seung-Woo Seo (Department of Electrical and Computer Engineering, Seoul National University) and Sang-Hyun Lee (Department of Automotive Engineering, Ajou University, corresponding author), with the paper listed as a Machine Learning / ICML submission and an equal-advising note.

What This Paper Is About

Offline goal-conditioned RL agents must learn to reach specified goals using only a fixed, pre-collected dataset, but performance degrades as task horizons grow because discounting shrinks the gap between good and bad actions and Bellman errors compound. Existing hierarchical fixes split the problem into a high-level subgoal selector and a low-level controller, yet they typically emit only one intermediate subgoal, operate as separate networks, and lose sight of the final goal when the high-level target is wrong. CoGHP's goal is to replace that two-network, single-subgoal structure with one model that generates a whole sequence of intermediate subgoals before acting.

Key Contributions

1. Chain-of-thought reformulation for offline hierarchical RL. The paper adapts the chain-of-thought reasoning paradigm, in which a complex problem is decomposed into a sequence of intermediate steps, to offline hierarchical RL, recasting hierarchical decision-making as autoregressive generation of multiple latent subgoals that act as reasoning steps.
2. MLP-Mixer as a hierarchical RL backbone. The authors introduce the MLP-Mixer architecture, augmented with a learnable causal token-mixer, as the backbone for hierarchical RL, using token-mixing and channel-mixing MLPs for cross-token communication between state, goal, latent subgoals and action, and enabling unified end-to-end training across all decision-making stages.
3. A single shared-value training objective spanning all hierarchy levels. An AWR-style objective with a shared goal-conditioned value function provides training signals for both intermediate subgoals and primitive actions, so corrective gradient signals propagate across the whole hierarchy rather than being confined to separate modules.
4. Empirical validation on long-horizon navigation and manipulation. CoGHP is evaluated on the OGBench suite against eight baselines, with ablation studies isolating the contributions of the MLP-Mixer backbone and the causal mixer, a study of subgoal count, and visualizations of the generated subgoal chains.

Main Findings

• Strongest results on the hardest mazes. On pointmaze-giant-navigate CoGHP reaches 79% versus HIQL's 46%; on antmaze-giant-navigate it reaches 78% versus HIQL's 65%. The authors attribute this to HIQL's two-level structure with separate networks being inadequate when multiple intermediate decisions must be coordinated.
• Large gains on sequential manipulation. On scene-play, which requires coordinating up to eight sequential atomic behaviors in the correct order, CoGHP achieves 78% compared with HIQL's 38%, OTA's 20% and SAW's 63%.
• Reversing HIQL's failure mode on cube placement. HIQL drops to 41% on cube-single while GCIQL reaches 99%, which the authors attribute to the low-level policy losing awareness of the final goal. CoGHP reaches 97% on cube-single and 42% on cube-triple, outperforming OTA and SAW by a large margin on multi-object manipulation.
• Competitive but not always best on easier tasks. On antmaze-medium-navigate, antmaze-large-navigate and pointmaze-medium-navigate the strongest baselines (SAW, OTA, GCIQL, QRL) match or slightly trail CoGHP, so the advantage is concentrated in the most complex environments.
• MLP-Mixer beats Transformer, and the gap grows with complexity. In the architecture ablation, both backbones tie at 97% on antmaze-medium-navigate, but on cube-triple the Transformer variant scores 2% versus CoGHP's 42%, and on antmaze-giant-navigate 66% versus 78%. The authors state that Transformer backbones offered limited generalization benefits and reduced training stability when tokens have fixed position-dependent roles.
• The causal mixer matters most where hierarchical reasoning is hardest. Removing it changes nothing on antmaze-medium-navigate and cube-single, but on cube-triple performance falls from 42% to 27% and on cube-double from 54% to 44%.
• More subgoals is not monotonically better. The optimal subgoal count is task-dependent: H=0 performs poorly on antmaze-large and antmaze-giant, antmaze-giant performs best with H=2, antmaze-large degrades with H=10, and both cube-double and scene achieve their best performance with H=1, with larger counts consistently performing worse.
• Generated subgoals lie on the optimal path. Visualizing three decoded subgoals in antmaze-giant, all three lie on or near the optimal path, and because generation is autoregressive the subgoal nearest the agent is conditioned on the previously generated ones.

Methodology in Plain English

The setting is an offline MDP with state space, action space, transition dynamics, a goal-dependent reward, a discount factor and an initial state distribution; the agent gets a fixed dataset of trajectories and no further environment interaction, and the goal space is assumed to equal the state space.

The core idea is to stop building a separate network for every level of the hierarchy. Instead, the model receives the current state and the final goal as embeddings, followed by placeholder tokens for H latent subgoals and one for the action. It then fills those tokens in one at a time. Crucially, it generates the subgoals in reverse order of distance, starting with the one farthest from the current state and ending with the one nearest, based on the hypothesis that subgoals closer to the agent should absorb the most information from the reasoning chain. The final goal stays present as a constant condition for every prediction, so the model cannot simply forget where it is ultimately heading. After all subgoals are produced, the action is emitted from the action token's hidden state.

The backbone is an MLP-Mixer, which alternates token-mixing MLPs (exchange information across tokens) with channel-mixing MLPs (refine features within each token). Since Mixer layers are inherently order-sensitive and need no separate positional embeddings, the authors add a learnable causal token-mixer implemented as a lower-triangular matrix, so each token becomes a weighted sum of previously generated tokens. The authors motivate this over a Transformer by arguing that in this setting each token has a fixed semantic role, so the dynamic attention of a Transformer yields little benefit and costs stability.

Training uses a shared goal-conditioned state-value function learned with the IQL temporal-difference objective, with the goal embedded so the same value function can score both embedded goals and latent subgoals, which live in the same latent space. Two AWR-style advantage-weighted losses are defined: one for each subgoal generation step, using the advantage of reaching the intermediate state relative to the goal, and one for the action, using the advantage of the action relative to the nearest generated subgoal. These are combined into a single end-to-end loss with weight coefficients for the subgoal and action terms and a discount that down-weights distant subgoals. Target subgoals are obtained by sampling states at fixed k-step intervals along dataset trajectories. Value function and policy are optimized alternately, and teacher forcing supplies ground-truth subgoal embeddings during policy training to prevent error accumulation early on. Latent subgoals are implemented simply as encoded future states.

Evaluation uses OGBench: pointmaze (a 2-D point mass) and antmaze (a quadrupedal Ant with 8 degrees of freedom) at medium, large and giant sizes, plus cube single/double/triple variants and the scene environment. Following OGBench's protocol, five predefined state-goal pairs per environment are tested and average success rate is reported, with standard deviations across 8 random seeds. Baselines are GCBC, GCIVL, GCIQL, QRL, CRL, HIQL, OTA and SAW. Pixel-based results (visual-antmaze, visual-cube), hyperparameter details and further ablation are deferred to appendices.

Why This Matters

Impact on research. The paper argues that the standard two-network, single-subgoal hierarchy is a structural bottleneck rather than a tuning problem, and offers an alternative that scales to multiple hierarchy levels without adding networks. It also opens a cross-pollination path: the chain-of-thought decomposition that has already been applied to vision-language-action models is ported into offline goal-conditioned RL and instantiated concretely. Finally, it is presented as the first adaptation of MLP-Mixer to offline goal-conditioned RL, which suggests that all-MLP backbones are worth revisiting in domains where token positions have fixed semantic roles.

Real-world applications (potential, as motivated by offline RL's premise of costly or unsafe interaction):

• Robot manipulation requiring ordered multi-step sequences, such as unlocking, opening, placing and closing, which is directly the structure of the scene environment.
• Long-horizon navigation for mobile robots, matching the pointmaze and antmaze settings.
• Training control policies from existing logged data in domains where online exploration is expensive or dangerous, which is the motivating premise of offline RL.
• Multi-object rearrangement and pick-and-place tasks, as in the single, double and triple cube variants.

Industry relevance. The pitch of offline RL is that policies can be learned from static, pre-collected datasets rather than live environment interaction, which matters when each trial is costly or unsafe. If multi-stage reasoning can be folded into one model instead of a stack of separately trained networks, that simplifies training pipelines and reduces the moving parts an engineering team must tune.

Future Directions

• Richer subgoal representations. The authors note that latent subgoals are currently implemented as encoded future states, and that learned skill primitives or abstract semantic embeddings could in principle be supported by the MLP-Mixer backbone, but would likely require additional data modalities or annotations plus dedicated training objectives. They explicitly leave this to future work.
• Automatic subgoal count selection. The subgoal count analysis shows the optimum is task-dependent, with antmaze-giant favouring H=2 and cube-double and scene favouring H=1, and larger counts hurting manipulation. No rule for choosing H is offered.
• Understanding the Transformer failure. The ablation shows Transformer backbones degrade sharply with complexity, but the paper defers further analysis of that baseline to an appendix, leaving the underlying cause to be pinned down.
• Broader evaluation. Pixel-based tasks appear only in an appendix, and the fairness of comparing against HIQL-style single-subgoal hierarchies could be probed further.

Target Audience

Researchers and graduate students working on offline reinforcement learning, goal-conditioned RL and hierarchical control, particularly those interested in long-horizon decision-making and in whether sequence-model backbones other than Transformers are suitable for RL. It also suits practitioners who need policies learned from fixed logged datasets for robotics or navigation, and readers following the transfer of chain-of-thought-style decomposition from language and vision-language-action models into control. Readers without background in value-based offline RL, advantage-weighted regression and hierarchical subgoal decomposition will find the method sections difficult without supplementary reading.

Authors’ abstract

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/