Small Dataset, Big Gains: Enhancing Reinforcement Learning by Offline Pre-Training with Model Based Augmentation

📄 arXiv: 2312.09844v2 📥 PDF

作者: Girolamo Macaluso, Alessandro Sestini, Andrew D. Bagdanov

分类: cs.LG, cs.AI

发布日期: 2023-12-15 (更新: 2023-12-19)

💡 一句话要点

提出基于模型增强的离线预训练方法，提升小数据集强化学习性能

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 离线强化学习 数据增强 世界模型 预训练 机器人控制

📋 核心要点

  1. 离线强化学习受限于数据集质量和大小，直接预训练可能导致次优策略。
  2. 利用世界模型进行数据增强，扩展离线数据集，提升预训练策略质量。
  3. 在MuJoCo任务上验证，显著减少在线微调所需的环境交互次数。

📝 摘要（中文）

离线强化学习利用预先收集的转移数据集来训练策略，可以作为在线算法的有效初始化，提高样本效率并加速收敛。然而，当这些数据集的规模和质量有限时，离线预训练可能会产生次优策略，并导致在线强化学习性能下降。本文提出了一种基于模型的数据增强策略，以最大化离线强化学习预训练的优势，并减少其有效性所需的数据规模。我们的方法利用在离线数据集上训练的环境世界模型，在离线预训练期间增强状态。我们在各种 MuJoCo 机器人任务上评估了我们的方法，结果表明它可以快速启动在线微调，并大幅减少（在某些情况下达到一个数量级）所需的环境交互次数。

🔬 方法详解

问题定义：离线强化学习依赖预先收集的数据集，但小规模或低质量的数据集会导致预训练的策略表现不佳，进而影响在线微调的效果。现有方法难以充分利用有限的离线数据，导致样本效率低下。

核心思路：利用离线数据训练一个环境的世界模型，然后使用该模型生成更多的数据，从而增强离线数据集。通过在增强的数据集上进行预训练，可以获得更好的策略初始化，从而加速在线微调过程。

技术框架：该方法包含两个主要阶段：1) 世界模型训练：使用离线数据集训练一个能够预测环境状态转移的模型。2) 离线预训练与数据增强：在原始离线数据集的基础上，利用训练好的世界模型生成新的状态，并将这些状态添加到原始数据集中。然后，使用增强后的数据集进行离线策略预训练。最后，使用预训练的策略作为在线强化学习的初始化。

关键创新：核心创新在于利用世界模型进行数据增强，从而克服了小规模离线数据集的限制。与传统的离线强化学习方法相比，该方法能够更有效地利用有限的数据，并获得更好的策略初始化。

关键设计：世界模型的具体结构和训练方式未知，论文中可能未详细说明。数据增强的具体策略，例如生成多少新的状态，以及如何选择生成的状态，也是关键的设计细节。损失函数的设计需要保证世界模型能够准确地预测环境状态的转移。

📊 实验亮点

实验结果表明，该方法在各种MuJoCo机器人任务上能够显著提升在线强化学习的性能。与直接使用原始离线数据集进行预训练相比，该方法能够将所需的在线交互次数减少一个数量级。这表明该方法能够有效地利用有限的离线数据，并获得更好的策略初始化。

🎯 应用场景

该研究成果可应用于机器人控制、自动驾驶等领域，尤其是在难以获取大量真实环境交互数据的场景下。通过离线预训练和数据增强，可以显著降低在线学习的成本，加速策略的收敛，并提高最终性能。该方法还有潜力应用于其他强化学习任务，例如游戏AI和推荐系统。

📄 摘要（原文）

Offline reinforcement learning leverages pre-collected datasets of transitions to train policies. It can serve as effective initialization for online algorithms, enhancing sample efficiency and speeding up convergence. However, when such datasets are limited in size and quality, offline pre-training can produce sub-optimal policies and lead to degraded online reinforcement learning performance. In this paper we propose a model-based data augmentation strategy to maximize the benefits of offline reinforcement learning pre-training and reduce the scale of data needed to be effective. Our approach leverages a world model of the environment trained on the offline dataset to augment states during offline pre-training. We evaluate our approach on a variety of MuJoCo robotic tasks and our results show it can jump-start online fine-tuning and substantially reduce - in some cases by an order of magnitude - the required number of environment interactions.