DI-engine中GuidedCostRewardModel概率张量维度问题的分析与解决

在强化学习框架DI-engine的逆向强化学习实现中，GuidedCostRewardModel模块负责从专家示范数据中学习奖励函数。近期在实际应用中发现该模块在处理连续动作空间时存在概率张量维度不匹配的问题，导致模型训练过程中出现运行时错误。本文将从技术角度深入分析问题成因，并提出工程解决方案。

问题现象

当使用CALVIN数据集进行训练时，系统抛出RuntimeError异常，提示在调用default_collate函数时发现张量维度不匹配。具体表现为：

• 专家数据中的概率张量维度为[1]
• 采样数据中的概率张量维度为[1,1] 这种维度差异导致无法正常执行张量堆叠操作。

技术背景

GuidedCostRewardModel的核心算法基于最大熵逆向强化学习框架，其训练过程需要同时处理两类数据：

1. 专家示范数据：标记概率固定为1
2. 策略采样数据：通过当前策略计算动作概率

在连续动作空间场景下，系统使用高斯分布建模策略输出，通过概率密度函数计算特定动作的似然值。原始实现中对概率值的维度处理存在不一致性。

问题根源分析

深入代码实现后发现维度差异主要来源于：

1. 专家数据处理时直接创建标量概率值（shape=[1]）
2. 采样数据处理时对log_prob执行unsqueeze(0)操作（shape=[1,1]）
3. 两种数据混合后调用default_collate时维度检查失败

解决方案

通过以下修改确保维度一致性：

1. 移除采样数据中不必要的unsqueeze操作
2. 在损失计算时统一添加维度扩展
3. 显式处理设备转移确保计算图完整性

关键修改点包括：

# 原始代码
samp[i]['prob'] = torch.exp(log_prob).unsqueeze(0)

# 修改后
samp[i]['prob'] = torch.exp(log_prob)
prob = samp['prob'].unsqueeze(-1)  # 延迟维度处理

工程实践建议

1. 维度一致性检查：在混合不同来源数据前应验证关键张量的shape
2. 设备管理：显式处理跨设备数据转移，避免隐式转换
3. 延迟变形：保持原始数据维度，在最终计算时统一处理
4. 单元测试：增加对collate函数的专项测试用例

总结

该问题的解决不仅修复了运行时错误，更体现了强化学习系统实现中的几个重要原则：

1. 数据预处理阶段应保持原始形态
2. 维度操作应当显式且有明确目的
3. 设备管理需要系统化处理

这种解决方案既保持了算法数学表达的正确性，又满足了工程实现的鲁棒性要求，为DI-engine中逆向强化学习组件的稳定运行提供了保障。