Partial Policy Gradients for RL in LLMs
This work addresses the problem of learning more reliably in reinforcement learning for LLMs by simplifying policy structures, which could benefit researchers working on conversational AI.
This paper proposes a method to optimize for a subset of future rewards in policy gradients, allowing for the modeling and comparison of different policy classes such as full planning, greedy, K-step lookahead, and segment policies. They evaluate these policies on persona-alignment conversational problems, finding that different policies excel in different problems.
Reinforcement learning is a framework for learning to act sequentially in an unknown environment. We propose a natural approach for modeling policy structure in policy gradients. The key idea is to optimize for a subset of future rewards: smaller subsets represent simpler policies, which can be learned more reliably because their empirical gradient estimates are more accurate. Our approach allows for modeling and comparison of different policy classes, including full planning, greedy, K-step lookahead, and segment policies. We evaluate the policies empirically on multiple persona-alignment conversational problems. Different policies excel in different problems, reflecting their different characteristics and highlighting the importance of our studied policy class.