GEM: Guided Expectation-Maximization for Behavior-Normalized Candidate Action Selection in Offline RL
This addresses a specific bottleneck in offline RL deployment for practitioners, but it is incremental as it builds on existing methods like Gaussian Mixture Models and candidate-based selection.
The paper tackles the problem of unreliable action selection in offline reinforcement learning when datasets have multimodal action distributions, by introducing GEM, a framework that uses Gaussian Mixture Models for multimodal and controllable action selection with behavior-normalized support, achieving competitive performance on D4RL benchmarks and offering a tunable compute-quality trade-off via candidate count.
Offline reinforcement learning (RL) can fit strong value functions from fixed datasets, yet reliable deployment still hinges on the action selection interface used to query them. When the dataset induces a branched or multimodal action landscape, unimodal policy extraction can blur competing hypotheses and yield "in-between" actions that are weakly supported by data, making decisions brittle even with a strong critic. We introduce GEM (Guided Expectation-Maximization), an analytical framework that makes action selection both multimodal and explicitly controllable. GEM trains a Gaussian Mixture Model (GMM) actor via critic-guided, advantage-weighted EM-style updates that preserve distinct components while shifting probability mass toward high-value regions, and learns a tractable GMM behavior model to quantify support. During inference, GEM performs candidate-based selection: it generates a parallel candidate set and reranks actions using a conservative ensemble lower-confidence bound together with behavior-normalized support, where the behavior log-likelihood is standardized within each state's candidate set to yield stable, comparable control across states and candidate budgets. Empirically, GEM is competitive across D4RL benchmarks, and offers a simple inference-time budget knob (candidate count) that trades compute for decision quality without retraining.