让大模型通过少量示例快速学会决策,尤其擅长从差数据中学习。
Scalable In-Context Q-Learning

- 用多头结构同时预测策略和价值函数,提升推理效率。
- 在多种环境上优于基线,尤其在使用低质量数据时仍表现稳定。
- 适合需要快速适应新任务的强化学习场景,如机器人控制。
近期语言模型在上下文学习方面取得显著进展,推动了将此能力拓展至决策领域的上下文强化学习(ICRL)。由于涉及更复杂的动态和时间相关性,现有方法在利用次优轨迹和实现精确上下文推断方面面临挑战。本文提出可扩展的上下文Q学习(S-ICQL),通过动态规划与世界建模,实现高效奖励最大化与任务泛化,同时保持监督预训练的可扩展性和稳定性。设计基于提示的多头变压器架构,用独立头部同时预测最优策略和上下文价值函数。预训练通用世界模型以捕捉任务相关信息,构建紧凑提示,支持快速精准的上下文推断。训练中通过拟合状态值函数到Q函数的上期望分位数进行迭代策略改进,并使用优势加权回归将上下文价值函数提炼为策略提取。在多种离散与连续环境中进行大量实验,结果表明其性能持续优于各类基线,尤其在从次优数据学习时表现突出。代码已公开于:https://github.com/NJU-RL/SICQL。
原文摘要 · Abstract (English)
Recent advancements in language models have demonstrated remarkable in-context learning abilities, prompting the exploration of in-context reinforcement learning (ICRL) to extend the promise to decision domains. Due to involving more complex dynamics and temporal correlations, existing ICRL approaches may face challenges in learning from suboptimal trajectories and achieving precise in-context inference. In the paper, we propose \textbf{S}calable \textbf{I}n-\textbf{C}ontext \textbf{Q}-\textbf{L}earning (\textbf{S-ICQL}), an innovative framework that harnesses dynamic programming and world modeling to steer ICRL toward efficient reward maximization and task generalization, while retaining the scalability and stability of supervised pretraining. We design a prompt-based multi-head transformer architecture that simultaneously predicts optimal policies and in-context value functions using separate heads. We pretrain a generalized world model to capture task-relevant information, enabling the construction of a compact prompt that facilitates fast and precise in-context inference. During training, we perform iterative policy improvement by fitting a state value function to an upper-expectile of the Q-function, and distill the in-context value functions into policy extraction using advantage-weighted regression. Extensive experiments across a range of discrete and continuous environments show consistent performance gains over various types of baselines, especially when learning from suboptimal data. Our code is available at \textcolor{magenta}{\href{https://github.com/NJU-RL/SICQL}{https://github.com/NJU-RL/SICQL}}.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。