Learning Temporal Abstractions via Variational Homomorphisms in Option-Induced Abstract MDPs

TL;DR

This paper introduces a novel hierarchical reinforcement learning framework that learns latent temporal abstractions called options, enabling efficient implicit reasoning in complex tasks, with theoretical guarantees and practical applications in language and control.

Contribution

It develops VMOC, an off-policy variational algorithm for learning diverse options, extends MDP homomorphism theory for abstract reasoning, and proposes a supervised fine-tuning initialization method.

Findings

Achieves strong performance on logical reasoning benchmarks.

Effectively transfers to complex locomotion tasks.

Provides a principled theoretical foundation for abstract reasoning.

Abstract

Large Language Models (LLMs) have shown remarkable reasoning ability through explicit Chain-of-Thought (CoT) prompting, but generating these step-by-step textual explanations is computationally expensive and slow. To overcome this, we aim to develop a framework for efficient, implicit reasoning, where the model "thinks" in a latent space without generating explicit text for every step. We propose that these latent thoughts can be modeled as temporally-extended abstract actions, or options, within a hierarchical reinforcement learning framework. To effectively learn a diverse library of options as latent embeddings, we first introduce the Variational Markovian Option Critic (VMOC), an off-policy algorithm that uses variational inference within the HiT-MDP framework. To provide a rigorous foundation for using these options as an abstract reasoning space, we extend the theory of continuous MDP homomorphisms. This proves that learning a policy in the simplified, abstract latent space, for which VMOC is suited, preserves the optimality of the solution to the original, complex problem. Finally, we propose a cold-start procedure that leverages supervised fine-tuning (SFT) data to distill human reasoning demonstrations into this latent option space, providing a rich initialization for the model's reasoning capabilities. Extensive experiments demonstrate that our approach achieves strong performance on complex logical reasoning benchmarks and challenging locomotion tasks, validating our framework as a principled method for learning abstract skills for both language and control.