[论文解读] Conditional Mutual Information for Disentangled Representations in Reinforcement Learning

Mhairi Dunion, Trevor McInroe+3 more|VU Research Portal|May 23, 2023

Adversarial Robustness in Machine Learning被引用 5

一句话总结

该论文提出条件互信息解耦（CMID），一种辅助训练目标，通过最小化潜在特征之间的条件互信息，在强化学习中学习解耦表征，即使变化因素之间存在相关性也能实现。CMID在连续控制环境中使用图像观测时，平均使零样本泛化能力提升77%，并增强了在特征相关性发生变化时的训练性能。

ABSTRACT

Reinforcement Learning (RL) environments can produce training data with spurious correlations between features due to the amount of training data or its limited feature coverage. This can lead to RL agents encoding these misleading correlations in their latent representation, preventing the agent from generalising if the correlation changes within the environment or when deployed in the real world. Disentangled representations can improve robustness, but existing disentanglement techniques that minimise mutual information between features require independent features, thus they cannot disentangle correlated features. We propose an auxiliary task for RL algorithms that learns a disentangled representation of high-dimensional observations with correlated features by minimising the conditional mutual information between features in the representation. We demonstrate experimentally, using continuous control tasks, that our approach improves generalisation under correlation shifts, as well as improving the training performance of RL algorithms in the presence of correlated features.

研究动机与目标

• 解决强化学习环境中虚假相关性带来的泛化与鲁棒性挑战。
• 克服现有解耦方法假设变化因素独立的局限性，这些方法在因素相关时会失效。
• 开发一种在条件独立而非严格独立条件下实现解耦表征的方法。
• 在部署时特征相关性发生变化或缺失的情况下，提升强化学习智能体的零样本泛化能力和训练稳定性。
• 展示CMID作为即插即用的辅助任务，可与SVEA、DrQ、CURL和TED等现有强化学习算法兼容。

提出的方法

• CMID引入一种辅助损失，通过最小化给定条件集下潜在特征对之间的条件互信息（CMI）来实现。
• 条件集基于马尔可夫决策过程（MDP）的因果结构构建，利用最近的观测使特征条件独立。
• 该方法使用k近邻（kNN）估计CMI的边缘分布乘积，从而实现可微训练。
• 潜在表征通过反向传播更新，损失函数结合了基础强化学习算法损失与CMID辅助损失。
• 使用时间上堆叠的特征表示而非原始帧，以避免在输入中引入虚假的因果依赖。
• 该方法兼容在线强化学习算法，可无需修改主策略训练循环直接应用。

实验结果

研究问题

• RQ1 当变化因素相关而非独立时，能否学习到解耦表征？
• RQ2 最小化潜在特征之间的条件互信息是否能提升强化学习中在相关性变化下的零样本泛化能力？
• RQ3 CMID与SOTA自监督强化学习方法（如DrQ、CURL和TED）相比，在应对特征相关性变化时的鲁棒性如何？
• RQ4 CMID能否在特征相关环境中提升训练性能和样本效率？
• RQ5 CMID作为无需架构修改的辅助任务，应用于现有强化学习算法时是否有效？

主要发现

• 在未见物体颜色的测试中，CMID相比基础SVEA算法，平均使零样本泛化性能提升77%。
• 在CartPole摆动任务中，CMID在最佳未见颜色下的平均回报达到834.1 ± 105.8，而SVEA为588.7 ± 87.2，表明其对颜色相关性变化具有显著鲁棒性。
• CMID在相关性变化下的泛化性能优于SOTA基线方法DrQ、CURL和TED，尤其在特征相关性阻碍学习的任务中表现更优。
• CMID提升了SVEA的训练性能，表明其有助于在相关特征存在时实现更快收敛和更稳定的训练。
• CMID成功通过学习条件独立特征，实现了对相关因素（如物体颜色与大小）的解耦，即使这些因素共同预测动力学行为。
• CMID的计算成本因kNN-based CMI估计使运行时间平均增加67%，但这一代价被泛化能力和鲁棒性的显著提升所抵消。