PWM: 基于世界模型的策略学习 原创

强化学习（RL）在复杂任务上取得了令人瞩目的成果，但在具有不同实施方式的多任务设置中存在困难。世界模型通过学习环境的模拟来提供可伸缩性，但它们通常依赖于低效的无梯度优化方法。近日，佐治亚理工学者联合英伟达、加州大学圣地亚哥分校等学者提出了基于大世界模型的策略学习（PWM），这是一种新颖的基于模型的RL算法，它从大型多任务世界模型中学习连续控制策略。通过对离线数据进行世界模型的预训练，并将其用于一阶梯度策略学习，PWM有效地解决了具有多达152个动作维度的任务，并且胜过使用真实动力学的方法。此外，PWM在80个任务设置中进行了扩展，相比于现有基准方法，其奖励提高了高达27％，而无需昂贵的在线规划。

方法概述

PWM是一种新颖的基于模型的RL（MBRL）算法和框架，旨在从大型多任务世界模型中得出有效的连续控制策略。利用预训练的TD-MPC2世界模型，在每个任务中以小于10分钟的时间高效地学习控制策略。对复杂的运动任务进行的实证评估表明，PWM不仅在奖励上超过了基准方法，还胜过了使用真实模拟动力学的方法。

PWM的预告结果

对高维连续控制任务上的PWM进行评估（左图），发现它不仅胜过了无模型的基准方法SAC和PPO，还在奖励上超过了直接使用模拟器动力学和奖励函数的SHAC方法。在一个包含80个任务的设置中（右图），使用了一个庞大的4800万参数的世界模型，PWM能够持续胜过使用相同世界模型的TD-MPC2方法，而无需在线规划。

单任务结果

图中显示了50％的IQM（实线）、均值（虚线）和所有5个任务和5个随机种子的95％置信区间。PWM能够获得比无模型的基准方法PPO和SAC、使用与PWM相同世界模型的TD-MPC2以及使用真实动力学和奖励函数的SHAC更高的奖励。这些结果表明，经过良好正则化的世界模型可以平滑优化空间，从而实现更好的一阶梯度优化。

多任务结果

图中显示了PWM和TD-MPC2在30个和80个多任务基准测试上的性能，结果基于10个随机种子。PWM能够在使用相同世界模型且无需任何形式的在线规划的情况下胜过TD-MPC2，这使其成为大型世界模型更可扩展的方法。右图将PWM（多任务策略）与单任务专家SAC和DreamerV3进行了比较。令人印象深刻的是，PWM能够在多任务情况下与它们的性能相匹配，同时只使用离线数据进行训练。

译自（有删改）：https://www.imgeorgiev.com/pwm

本文转载自公众号AIGC最前线   

原文链接：​​https://mp.weixin.qq.com/s/UQqEdQl1YIuwqp-3ytCMpA​​