通过时序差分预测潜在状态,实现无需奖励的零样本强化学习。
TD-JEPA: Latent-predictive Representations for Zero-Shot Reinforcement Learning
- 基于时序差分构建多步潜在预测模型,从离线数据中学习通用表示。
- 在13个数据集上达到或超过现有最优性能,尤其在像素级零样本任务中表现突出。
- 适合追求无监督预训练与灵活奖励部署的研究者和应用开发者。
潜在预测——即智能体通过预测自身潜在表示进行学习——已成为机器学习中训练通用表示的强大范式。在强化学习(RL)中,该方法被用于多种场景的辅助损失设计,包括基于奖励、无监督RL、行为克隆和世界建模。尽管现有方法通常局限于单任务学习、单步预测或在线轨迹数据,我们证明时序差分(TD)学习可实现跨多策略的长期潜在动态预测,且仅需离线、无奖励的转移数据。在此基础上,我们提出TD-JEPA,将基于TD的潜在预测表示引入无监督强化学习。TD-JEPA同时训练显式状态与任务编码器、条件于策略的多步预测器以及一组参数化策略,直接在潜在空间中操作。这使得测试时可对任意奖励函数实现零样本优化。理论上,理想化的TD-JEPA变体在合理初始化下避免退化,学习到能捕捉长期策略动态低秩分解的编码器,同时预测器恢复了潜在空间中的后继特征。实验表明,TD-JEPA在ExoRL和OGBench共13个数据集上的运动、导航与操作任务中匹配或超越当前最优基线,尤其在像素级零样本强化学习这一挑战性设置中表现优异。
原文摘要 · Abstract (English)
Latent prediction--where agents learn by predicting their own latents--has emerged as a powerful paradigm for training general representations in machine learning. In reinforcement learning (RL), this approach has been explored to define auxiliary losses for a variety of settings, including reward-based and unsupervised RL, behavior cloning, and world modeling. While existing methods are typically limited to single-task learning, one-step prediction, or on-policy trajectory data, we show that temporal difference (TD) learning enables learning representations predictive of long-term latent dynamics across multiple policies from offline, reward-free transitions. Building on this, we introduce TD-JEPA, which leverages TD-based latent-predictive representations into unsupervised RL. TD-JEPA trains explicit state and task encoders, a policy-conditioned multi-step predictor, and a set of parameterized policies directly in latent space. This enables zero-shot optimization of any reward function at test time. Theoretically, we show that an idealized variant of TD-JEPA avoids collapse with proper initialization, and learns encoders that capture a low-rank factorization of long-term policy dynamics, while the predictor recovers their successor features in latent space. Empirically, TD-JEPA matches or outperforms state-of-the-art baselines on locomotion, navigation, and manipulation tasks across 13 datasets in ExoRL and OGBench, especially in the challenging setting of zero-shot RL from pixels.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。