GEM: Guided Expectation-Maximization for Behavior-Normalized Candidate Action Selection in Offline RL

TL;DR

GEM introduces a multimodal, controllable action selection framework for offline RL that improves decision robustness by combining candidate generation, behavior normalization, and ensemble reranking, adaptable to computational budgets.

Contribution

GEM presents a novel analytical framework that enables explicit control over multimodal action selection in offline RL through a candidate-based, behavior-normalized approach.

Findings

GEM performs competitively on D4RL benchmarks.

GEM allows adjustable inference complexity via candidate count.

GEM enhances decision robustness in multimodal action landscapes.

Abstract

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.