Improving and Accelerating Offline RL in Large Discrete Action Spaces with Structured Policy Initialization

TL;DR

This paper introduces SPIN, a two-stage framework for offline reinforcement learning in large discrete action spaces, which improves performance and training speed by leveraging structured policy initialization.

Contribution

The paper proposes a novel two-stage approach, SPIN, that pre-trains an action structure model and then trains lightweight policies, enhancing efficiency and effectiveness in large discrete action spaces.

Findings

SPIN improves average return by up to 39% over state-of-the-art methods.

SPIN reduces training time to convergence by up to 12.8 times.

Effective handling of combinatorial action spaces in offline RL.

Abstract

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$.