用流匹配建模智能体未来状态分布,通过意图条件提升强化学习预训练效果。
Intention-Conditioned Flow Occupancy Models
- 基于流匹配构建意图条件的未来状态分布模型,捕捉长期依赖。
- 在40个任务上平均收益提升1.8倍,成功率提高36%。
- 适合希望高效迁移的RL研究者,尤其关注通用策略改进场景。
大规模预训练已深刻改变机器学习研究范式:大型基础模型只需训练一次,即可被社区任何人(包括无数据或算力资源者)用于特定任务的适配与微调。将此框架应用于强化学习(RL)极具吸引力,因其有望解决样本效率和鲁棒性等核心挑战。然而,在RL中预训练大模型仍面临根本难题:动作具有长期依赖,需训练能跨时间推理的基础模型。生成式AI的进展为建模复杂分布提供了新工具。本文提出一种概率模型,利用流匹配预测智能体在遥远未来的状态访问情况(即占用度量)。由于大规模数据集常由不同用户执行不同任务构建,模型引入隐变量以捕捉用户意图。该意图显著提升模型表达能力,并支持广义策略改进。我们称该方法为意图条件流占用模型(InFOM)。在36个状态基和4个图像基基准任务上的实验表明,相比其他预训练方法,本方法在回报上实现1.8倍的中位数提升,成功率增加36%。
原文摘要 · Abstract (English)
Large-scale pre-training has fundamentally changed how machine learning research is done today: large foundation models are trained once, and then can be used by anyone in the community (including those without data or compute resources to train a model from scratch) to adapt and fine-tune to specific tasks. Applying this same framework to reinforcement learning (RL) is appealing because it offers compelling avenues for addressing core challenges in RL, including sample efficiency and robustness. However, there remains a fundamental challenge to pre-train large models in the context of RL: actions have long-term dependencies, so training a foundation model that reasons across time is important. Recent advances in generative AI have provided new tools for modeling highly complex distributions. In this paper, we build a probabilistic model to predict which states an agent will visit in the temporally distant future (i.e., an occupancy measure) using flow matching. As large datasets are often constructed by many distinct users performing distinct tasks, we include in our model a latent variable capturing the user intention. This intention increases the expressivity of our model, and enables adaptation with generalized policy improvement. We call our proposed method intention-conditioned flow occupancy models (InFOM). Comparing with alternative methods for pre-training, our experiments on $36$ state-based and $4$ image-based benchmark tasks demonstrate that the proposed method achieves $1.8 \times$ median improvement in returns and increases success rates by $36\%$. Website: https://chongyi-zheng.github.io/infom Code: https://github.com/chongyi-zheng/infom
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。