Improving and Accelerating Offline RL in Large Discrete Action Spaces with Structured Policy Initialization
This addresses the challenge of slow and unstable policy learning in combinatorial action spaces for offline RL, offering significant performance gains.
The paper tackles the problem of reinforcement learning in large discrete action spaces by introducing Structured Policy Initialization (SPIN), which pre-trains an Action Structure Model to capture valid actions and then trains lightweight policy heads, improving average return by up to 39% and reducing convergence time by up to 12.8x on DM Control benchmarks.
Reinforcement learning in discrete combinatorial action spaces requires searching over exponentially many joint actions to simultaneously select multiple sub-actions that form coherent combinations. Existing approaches either simplify policy learning by assuming independence across sub-actions, which often yields incoherent or invalid actions, or attempt to learn action structure and control jointly, which is slow and unstable. We introduce Structured Policy Initialization (SPIN), a two-stage framework that first pre-trains an Action Structure Model (ASM) to capture the manifold of valid actions, then freezes this representation and trains lightweight policy heads for control. On challenging discrete DM Control benchmarks, SPIN improves average return by up to 39% over the state of the art while reducing time to convergence by up to 12.8$\times$.