Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

Matthew Landers, Taylor W. Killian, Thomas Hartvigsen, Afsaneh Doryab

TL;DR

SPIN tackles offline RL in large discrete combinatorial action spaces by decoupling learning of action structure from control. It first trains an Action Structure Model (ASM) with masked conditional modeling to capture a low-dimensional, coherent action manifold, then freezes this representation and learns lightweight, context-aware policy heads over the manifold. Empirically, SPIN achieves state-of-the-art performance on challenging DM Control benchmarks, with up to a 39% improvement in average return and up to 12.8x faster convergence than strong baselines, and demonstrates robustness to increasing action cardinality. Analyses show that the gains stem from the quality of the learned action representation, with linear probes revealing substantial cross-joint coordination, and SPIN-Distill showing that a smaller policy can still reap the benefits when the representation is strong. Overall, SPIN provides a representation-first framework for scalable, efficient offline RL in structured action spaces, with potential extensions to alternative offline objectives and action-structure assumptions.

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

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

TL;DR

SPIN tackles offline RL in large discrete combinatorial action spaces by decoupling learning of action structure from control. It first trains an Action Structure Model (ASM) with masked conditional modeling to capture a low-dimensional, coherent action manifold, then freezes this representation and learns lightweight, context-aware policy heads over the manifold. Empirically, SPIN achieves state-of-the-art performance on challenging DM Control benchmarks, with up to a 39% improvement in average return and up to 12.8x faster convergence than strong baselines, and demonstrates robustness to increasing action cardinality. Analyses show that the gains stem from the quality of the learned action representation, with linear probes revealing substantial cross-joint coordination, and SPIN-Distill showing that a smaller policy can still reap the benefits when the representation is strong. Overall, SPIN provides a representation-first framework for scalable, efficient offline RL in structured action spaces, with potential extensions to alternative offline objectives and action-structure assumptions.

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.
Paper Structure (36 sections, 10 equations, 3 figures, 15 tables, 2 algorithms)

This paper contains 36 sections, 10 equations, 3 figures, 15 tables, 2 algorithms.

Table of Contents

  1. Introduction
  2. Related Work
  3. RL in Large Discrete Action Spaces.
  4. Self-Supervised Pre-Training in RL.
  5. Preliminaries
  6. Combinatorial action space.
  7. Offline reinforcement learning.
  8. Structured Policy Initialization (SPIN)
  9. Action Structure Modeling (ASM)
  10. Policy Learning with a Frozen Representation
  11. Policy Architecture and Training
  12. Compatibility with Offline RL Methods.
  13. Experimental Evaluation
  14. Asymptotic Performance and Training Efficiency
  15. Robustness to Action Cardinality
  16. ...and 21 more sections

Figures (3)

  • Figure 1: In the humanoid-stand task, learning from the medium-expert dataset, SPIN reaches the target performance (dashed horizontal line) quickly after pre-training (dashed vertical line), while baselines require over 200 minutes of wall-clock training.
  • Figure 2: Final policy return as a function of ASM pre-training duration. Policies are initialized from frozen ASMs trained for different numbers of epochs. While pre-training is critical, just 20 epochs yield a representation that enables policies to surpass the fully trained F-IQL baseline in all tasks (dashed horizontal line).
  • Figure 3: Learning curves on all DM Control tasks and dataset qualities. Curves show mean return over 1M gradient steps, averaged across five seeds with shaded regions denoting $\pm$1 std.