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

Haoyu Wang; Jingcheng Wang; Shunyu Wu; Xinwei Xiao

cs.LGMar 24, 2026

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

Authors: Haoyu Wang, Jingcheng Wang, Shunyu Wu, Xinwei Xiao

Organizations: School of Automation and Intelligent Sensing, Shanghai Jiao Tong University, Shanghai, China

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.

Figures & tables

Figure 1: GEM mechanism (schematic at a fixed state in a 2-D action space). Training (top row): GEM learns three ingredients that will be combined only at test time: (i) a multimodal actor $\pi_{\theta}(a\mid s)$ (left; shown as $\log\pi_{\theta}(a\mid s)$ ) that preserves multiple plausible action hypotheses, (ii) an independent behavior density $\mu_{\varphi}(a\mid s)$ (middle; shown as $\log\mu_{\varphi}(a\mid s)$ ) used to quantify dataset support, and (iii) a conservative value statistic $\mathrm{LCB}_{\lambda}(s,a)$ from a critic ensemble (right). Inference (bottom row): given a queried state $s$ , GEM samples $N$ candidates in parallel and adds a deterministic anchor (triangle), scores each candidate using Eq. 4 (conservative value plus behavior-normalized support), and executes the top-ranked action $a^{\star}$ (star). The green cluster highlights the small subset of candidates that remain competitive after scoring, illustrating how the interface filters many proposals down to a support-aligned, high-value choice under maximization.

Locomotion
Task	TD3+BC	AWAC	CQL	IQL	ReBRAC	SAC-N	EDAC	DT	IDQL	Diffusion-QL	FQL	SSAR	GEM(OURS)
HC-M	48.10	50.02	47.04	48.31	64.04	68.20	67.70	42.20	51.10	51.10	nan	60.00	54.42
HC-MR	44.84	45.13	45.04	44.46	51.18	60.70	62.06	38.91	45.90	47.80	nan	51.70	49.60
HC-ME	90.78	95.00	95.63	94.74	103.80	98.96	104.76	91.55	95.90	96.80	nan	98.50	100.95
H-M	60.30	63.02	59.08	67.53	102.29	40.82	101.70	65.10	65.40	90.50	nan	95.40	100.08
H-MR	64.42	98.88	95.11	97.43	94.98	100.33	99.66	81.77	92.10	101.30	nan	101.50	101.23

Table 1: Unified D4RL benchmark comparison (normalized score, higher is better). Suites: Locomotion (HalfCheetah/Hopper/Walker2d), AntMaze, and Maze2D.

Appendix figures & tables26 assets

Supplementary material from the paper’s appendix.

Appendix

Environment	Mean gap $\pm$ std
Maze2D-umaze	168.29 $\pm$ 7.71
Maze2D-medium	147.30 $\pm$ 6.85
Maze2D-large	142.64 $\pm$ 7.34
Antmaze-umaze-diverse	38.80 $\pm$ 3.46
Antmaze-umaze	31.53 $\pm$ 2.83
Antmaze-medium-diverse	23.73 $\pm$ 2.68

Appendix

Table 2: Per-environment NLL-gap statistics (

\mathrm{gap}=\mathrm{NLL}_{\mathrm{top1}}-\mathrm{NLL}_{\mathrm{gmm}}

; mean

\pm

std). All entries are positive in the provided results, indicating consistent mixture benefits.

Figure 5: Locomotion suite: per-environment candidate-budget scaling for Score.

Figure 6: Locomotion suite: per-environment candidate-budget scaling for Violation .

Figure 7: Locomotion suite: per-environment candidate-budget scaling for CollapseDist .

Figure 8: AntMaze suite: per-environment candidate-budget scaling for Score.

Figure 9: AntMaze suite: per-environment candidate-budget scaling for Violation .

Figure 10: AntMaze suite: per-environment candidate-budget scaling for CollapseDist .

Figure 11: Maze2D suite: per-environment candidate-budget scaling for Score.

Figure 12: Maze2D suite: per-environment candidate-budget scaling for Violation .

Figure 13: Maze2D suite: per-environment candidate-budget scaling for CollapseDist .

Figure 14: Robustness (aggregated) under mass-scaling perturbations on the Locomotion suite.

Figure 15: Robustness (aggregated) under mass-scaling perturbations on Maze2D.

Figure 16: Robustness (aggregated) under mass-scaling perturbations on AntMaze. Near-zero scores across MassScale indicate widespread failure under this dynamics shift under our evaluation protocol.

Panel of Figure 17 — Figure 17: Suite-level interface ablations at $N{=}1024$ (deltas relative to GEM). Top: $\Delta$ Score. Middle: $\Delta$ Violation . Bottom: $\Delta$ CollapseDist .

Figure 18: Effect of support normalization at $N{=}1024$ (Mode= zscore minus Mode= raw ). Top: $\Delta$ Score (positive is better). Middle: $\Delta$ Violation . Bottom: $\Delta$ CollapseDist . This comparison isolates scale sensitivity: raw $\log\mu_{\varphi}(a\mid s)$ can vary substantially across states, so its additive influence can make a single global $w_{p}$ behave inconsistently across the state space.

Method	Setting	Latency (ms)	Peak alloc (MB)	Peak reserved (MB)	Params (M)
Decision Transformer	N1	0.88 $\pm$ 0.05	20.71 $\pm$ 0.00	25.00 $\pm$ 0.00	5.27
Decision Transformer	N64	0.88 $\pm$ 0.05	20.71 $\pm$ 0.00	25.00 $\pm$ 0.00	5.27
Decision Transformer	N1024	0.88 $\pm$ 0.05	20.71 $\pm$ 0.00	25.00 $\pm$ 0.00	5.27
Decision Transformer	N2048	0.88 $\pm$ 0.05	20.71 $\pm$ 0.00	25.00 $\pm$ 0.00	5.27
Decision Transformer	T10	0.88 $\pm$ 0.05	20.71 $\pm$ 0.00	25.00 $\pm$ 0.00	5.27
Decision Transformer	T20	0.88 $\pm$ 0.05	20.71 $\pm$ 0.00	25.00 $\pm$ 0.00	5.27

Appendix

Table 3: Full deployment compute measurements (mean

\pm

std). Latency is in milliseconds; memory is peak allocated/reserved MB; Params is model parameter count (millions).

TestEnv	GEM	n	SupportZ	Viol	Collapse	SPS
antmaze-large-diverse-v2	1.00 $\pm$ 1.41	3	1.99	0.0019	0.308	2742
antmaze-large-play-v2	5.00 $\pm$ 1.41	3	2.02	0.0024	0.318	2760
antmaze-medium-diverse-v2	11.76 $\pm$ 10.54	3	1.86	0.0031	0.357	2913
antmaze-medium-play-v2	76.36 $\pm$ 33.43	3	1.05	0.0356	0.748	2962
antmaze-umaze-diverse-v2	57.37 $\pm$ 5.55	3	2.03	0.0002	0.278	3173
antmaze-umaze-v2	75.71 $\pm$ 20.55	3	0.81	0.0189	0.309	3195

Appendix

Table 4: GEM primary performance (Mode=zscore,

N=1024

) aggregated by D4RL task. Columns report normalized score (mean

\pm

std across training seeds), average support z-score, support violation rate, collapse distance, and evaluation throughput.

Variant	Mean	Delta	SupportZ	Viol	Collapse	SPS	n
GEM	76.88	0.00 $\pm$ 0.00	1.25	0.0069	0.273	3064	47
anchor_bary	76.29	-0.60 $\pm$ 3.87	1.24	0.0070	0.271	3043	47
behavior_only_cands	79.92	3.03 $\pm$ 20.58	0.93	0.0144	0.240	3077	47
no_anchor	77.04	0.15 $\pm$ 2.32	1.24	0.0069	0.273	3109	47
no_lcb	74.84	-2.04 $\pm$ 16.02	1.27	0.0085	0.270	3093	47
no_pw	65.31	-11.57 $\pm$ 37.56	-0.90	0.2169	0.647	3082	47

Appendix

Table 5: Mechanism ablations at

N=1024

(Mode=zscore), aggregated over all tasks and training seeds.

\Delta

is the per-task score difference relative to GEM (mean

\pm

std across tasks).

TestEnv	GEM	no-pw	no-lcb	unimodal-actor	behavior-only
antmaze-large-diverse-v2	0.00	+1.00	+0.00	+8.00	+0.00
antmaze-large-play-v2	4.00	+0.00	-1.00	-2.00	-4.00
antmaze-medium-diverse-v2	6.00	+68.14	+0.00	-4.00	+2.00
antmaze-medium-play-v2	76.47	-56.47	-37.05	+23.53	-19.80
antmaze-umaze-diverse-v2	52.42	-31.42	-9.06	-11.93	-16.23
antmaze-umaze-v2	71.58	-68.58	-5.00	+0.00	-6.00

Appendix

Table 6: Per-task ablation deltas at

N=1024

(Mode=zscore). Columns report mean normalized score of GEM and mean score change (

\Delta

) of each variant relative to GEM (positive is better).

N	Score	Std	SupportZ	Viol	Collapse	Steps/s	n
0	64.24	30.99	0.00	0.0000	0.265	5625	846
1	65.77	30.53	0.92	0.0000	0.271	5254	846
64	72.32	34.25	1.26	0.0027	0.257	5383	846
1024	72.79	34.46	1.33	0.0048	0.237	3105	846
2048	73.41	35.00	1.34	0.0055	0.233	1958	846

Appendix

Table 7: Test-time scaling with candidate budget

N

for GEM (Mode=zscore), aggregated over all sweep runs. Larger

N

improves average score while reducing throughput.

TestEnv	Delta	n
antmaze-large-diverse-v2	0.00 $\pm$ 0.00	2
antmaze-large-play-v2	5.00 $\pm$ 1.41	2
antmaze-medium-diverse-v2	9.76	1
antmaze-medium-play-v2	19.89 $\pm$ 25.51	2
antmaze-umaze-diverse-v2	3.81 $\pm$ 19.59	2
antmaze-umaze-v2	-1.12 $\pm$ 6.91	2

Appendix

Table 8: Effect of support normalization: baseline GEM difference between Mode=zscore and Mode=raw (reported as

\mathrm{Score}_{z}-\mathrm{Score}_{raw}

) aggregated by task.

Environment	Gap	Seeds
antmaze-large-diverse-v2	42.71 $\pm$ 7.04	3
antmaze-large-play-v2	43.22 $\pm$ 15.75	3
antmaze-medium-diverse-v2	57.16 $\pm$ 9.94	3
antmaze-medium-play-v2	48.45 $\pm$ 13.97	3
antmaze-umaze-diverse-v2	48.42 $\pm$ 7.08	3
antmaze-umaze-v2	48.86 $\pm$ 7.26	3

Appendix

Table 9: Multimodality diagnostic (dataset NLL gap). Larger gap indicates higher cost of a unimodal projection, motivating multimodal proposals.

Method	T10	T100	T20	T50
DT	0.934 $\pm$ 0.122	0.866 $\pm$ 0.008	0.874 $\pm$ 0.013	0.866 $\pm$ 0.010
FQL	1.479 $\pm$ 0.014	1.487 $\pm$ 0.018	1.480 $\pm$ 0.012	1.474 $\pm$ 0.024
GEM	1.296 $\pm$ 0.016	1.301 $\pm$ 0.043	1.270 $\pm$ 0.006	1.292 $\pm$ 0.052
IDQL	2.042 $\pm$ 0.017	20.367 $\pm$ 0.163	4.112 $\pm$ 0.014	10.096 $\pm$ 0.050
IFQL	1.831 $\pm$ 0.017	1.827 $\pm$ 0.010	1.821 $\pm$ 0.019	1.804 $\pm$ 0.026

Appendix

Table 10: Inference latency (ms) under diffusion-step sweep

T

(mean

\pm

std).

Method	N1	N1024	N2048	N64
DT	0.886 $\pm$ 0.006	0.884 $\pm$ 0.007	0.878 $\pm$ 0.006	0.865 $\pm$ 0.008
FQL	1.483 $\pm$ 0.015	1.506 $\pm$ 0.014	1.465 $\pm$ 0.038	1.482 $\pm$ 0.010
GEM	1.307 $\pm$ 0.007	1.276 $\pm$ 0.016	1.254 $\pm$ 0.002	1.279 $\pm$ 0.010
IDQL	10.168 $\pm$ 0.046	10.175 $\pm$ 0.051	10.115 $\pm$ 0.108	10.107 $\pm$ 0.072
IFQL	1.796 $\pm$ 0.008	1.837 $\pm$ 0.035	1.865 $\pm$ 0.076	1.830 $\pm$ 0.012

Appendix

Table 11: Inference latency (ms) under candidate-budget sweep

N

(mean

\pm

std).

Method	T10	T100	T20	T50
DT	21.1 $\pm$ 0.0	21.1 $\pm$ 0.0	21.1 $\pm$ 0.0	21.1 $\pm$ 0.0
FQL	20.7 $\pm$ 0.0	20.7 $\pm$ 0.0	20.7 $\pm$ 0.0	20.7 $\pm$ 0.0
GEM	50.8 $\pm$ 0.0	50.8 $\pm$ 0.0	50.8 $\pm$ 0.0	50.8 $\pm$ 0.0
IDQL	20.7 $\pm$ 0.0	20.7 $\pm$ 0.0	20.7 $\pm$ 0.0	20.7 $\pm$ 0.0
IFQL	20.8 $\pm$ 0.0	20.8 $\pm$ 0.0	20.8 $\pm$ 0.0	20.8 $\pm$ 0.0

Appendix

Table 12: Peak GPU allocated memory (MB) under diffusion-step sweep

T

(mean

\pm

std).

Method	N1	N1024	N2048	N64
DT	21.1 $\pm$ 0.0	21.1 $\pm$ 0.0	21.1 $\pm$ 0.0	21.1 $\pm$ 0.0
FQL	20.7 $\pm$ 0.0	20.7 $\pm$ 0.0	20.7 $\pm$ 0.0	20.7 $\pm$ 0.0
GEM	50.8 $\pm$ 0.0	50.8 $\pm$ 0.0	50.8 $\pm$ 0.0	50.8 $\pm$ 0.0
IDQL	20.7 $\pm$ 0.0	20.7 $\pm$ 0.0	20.7 $\pm$ 0.0	20.7 $\pm$ 0.0
IFQL	20.8 $\pm$ 0.0	20.8 $\pm$ 0.0	20.8 $\pm$ 0.0	20.8 $\pm$ 0.0

Appendix

Table 13: Peak GPU allocated memory (MB) under candidate-budget sweep

N

(mean

\pm

std).

Explore similar work

Decoupling Policy Extraction for Offline Reinforcement Learning
Aug 21, 2026Xuyao Lin, Yixiang Shan, Jinru Duan +7Offline RL Policy Learning
Simplexity Robotics, Rensselaer Polytechnic Institute · Simplexity Robotics · Simplexity Robotics, Northeastern University
Preserve Support, Not Correspondence: Dynamic Routing for Offline Reinforcement Learning
Apr 24, 2026Zhancun Mu, Guangyu Zhao, Yiwu Zhong +1Reinforcement Learning Offline RL
School of Intelligence Science and Technology, Peking University
ReBRAC-v2: The Return of the King
Aug 2, 2026Denis Tarasov, Robert K. KatzschmannReinforcement Learning Offline RL
Soft Robotics Lab, D-MAVT ETH Zurich, Switzerland

Aug 21, 2026J/K move · Enter open · S save

Decoupling Policy Extraction for Offline Reinforcement Learning

Xuyao Lin, Yixiang Shan, Jinru Duan, Tao Yang+6

Simplexity Robotics, Rensselaer Polytechnic Institute · Simplexity Robotics · Simplexity Robotics, Northeastern University

Offline RL Policy Learning

Open paper AskPDF arXiv

Offline RL methods commonly jointly train the actor and critic, where the critic is used to guide the actor toward higher-value actions. This coupled learning process is well motivated in online RL, where an improved actor collects new data that can further update the actor and the critic. However, training data remains fixed in offline RL, making actor-side policy improvement unable to generate new data to validate or correct the critic. Moreover, retaining this coupled paradigm leads to two related challenges. Firstly, actor updates can drift toward high-valued but potentially out-of-distribution (OOD) actions and amplify critic overestimation. Secondly, conservative value estimation or behavior-cloning regularization creates a difficult trade-off between suppressing OOD actions and selecting high-value actions within the data-supported region. Motivated by this observation, we revisit the conventional offline RL paradigm and propose decoupling policy improvement from actor training. Specifically, we train the actor solely to model the behavior distribution and perform policy improvement at inference time by reranking multiple actor-generated proposals with a separately learned critic. We refer to this paradigm as the decoupled policy extraction paradigm. Under such paradigm, the actor provides behavior-supported action candidates, while the critic performs value-based selection within this candidate set. Extensive experiments show that the decoupled policy extraction paradigm outperforms both behavior cloning and jointly learned offline RL methods, while remaining effective even with a naive Q-learning critic.

Abstract

Explore similar work

Decoupling Policy Extraction for Offline Reinforcement Learning

Preserve Support, Not Correspondence: Dynamic Routing for Offline Reinforcement Learning

ReBRAC-v2: The Return of the King