Energy-Based Policy Constraint for Offline Reinforcement Learning

Offline RL suffers from the distribution shift problem. One way to address this issue is to constrain the divergence between the target policy and the behavior policy. However, directly using the behavior policy-based constraint has two drawbacks: first, it does not directly distinguish the in-distribution samples and OOD samples, possibly overly constraining the target policy and limiting performance. Second, the practical datasets may be collected from multiple different behavior policies, which results in a multi-modal distribution, making it hard to represent the behavior policy. To address this problem, we propose a policy constraint method based on the energy-based model. On the one hand, the energy-based model constrains the target policy by energy function rather than directly constraining it to the dataset actions, making it suitable for multi-modal distribution. On the other hand, the energy-based model can effectively detect OOD samples, avoiding over-constraint of the target policy and improving the ceiling of the algorithm’s performance. The proposed algorithm is evaluated on the D4RL datasets. Experimental results show that compared to the behavior policy-based constraint methods, the energy-based policy constant significantly improves the performance and outperforms existing offline RL baselines.