Breaking the Performance Ceiling in Reinforcement Learning requires Inference Strategies¶

🕒 Published (v1): 2025-01-01 · Source: NeurIPS · Venue: NeurIPS 2025 · link

Ask a follow-up

Open an assistant pre-loaded with this paper's context.

TL;DR¶

Training RL policies to convergence is insufficient for complex multi-agent tasks — a dedicated inference phase with an appropriate search strategy yields up to 126% and on average 45% improvement over zero-shot SOTA across 17 hard tasks, using only seconds of extra wall-clock time. The paper formalizes inference strategies under a unified Dec-POMDP framework, extends four paradigms to the multi-agent setting, and provides the largest empirical study (60k+ experiments) of inference-time compute scaling for RL to date.

Problem¶

State-of-the-art MARL policies trained to convergence hit a hard performance ceiling on combinatorial, partially observable tasks (14/17 benchmark tasks remain below 70% normalized performance even after convergence). Existing evaluations exclusively measure zero-shot performance and neglect the inference phase, causing practitioners to expend effort chasing marginal training gains instead of leveraging inference-time search that is already feasible within seconds.

Method¶

The paper frames the problem as a Dec-POMDP (decentralized partially observable MDP) with a two-phase setup: unconstrained training followed by a budget-constrained inference phase I: (ρ, πθ, B_max, T_max) → (a*_1,...,a*_H). Four inference strategies are unified and extended to MARL:

Stochastic policy sampling — batch-parallel stochastic rollouts from the joint policy instead of a single greedy decode.
Simulation-guided beam search (SGBS) — expand top-K joint actions, simulate greedy rollouts, prune to best nodes; adapted for auto-regressive action selection via stochastic K-sampling from each node.
Online fine-tuning — gradient updates to policy parameters during inference using rollouts on the fixed instance, re-using the pre-training MARL algorithm.
COMPASS — a diversity-based method encoding specialized policies in a shared continuous latent space z ∈ L; at inference, CMA-ES searches L for the best latent vector. A single shared latent space prevents exponential growth with agent count.

All methods are implemented in JAX and evaluated on IPPO, MAPPO, and SABLE base policies across RWARE, SMACv2, and Connector environments under a grid of time (0–300 s) and compute (4–256 parallel attempts) budgets.

Key Contributions¶

Unified formal treatment of inference strategies in the multi-agent RL setting under a Dec-POMDP budget model.
Extension of stochastic sampling, SGBS, online fine-tuning, and COMPASS to the multi-agent setting.
Largest empirical study of inference strategies for RL: 60,928 evaluated episodes across 17 hard tasks, 4 strategies, 7 compute budgets, 128 seeds per task.
Contour-plot analysis of performance over (time, compute) budget grids, enabling practitioners to select strategies given their constraints.
Demonstration of compute scaling: SABLE+COMPASS shows near-linear performance scaling with additional parallel compute before saturating at 95%+ win-rate.

Results¶

Average 45% improvement over zero-shot SOTA across all 17 tasks with a 30-second / 64-parallel-attempt budget.
Up to 126% improvement on the hardest individual tasks.
SABLE+COMPASS achieves close to 100% win-rate on all tasks where win-rate is the metric, with 30 s of search.
Performance gains scale quasi-exponentially with task complexity — harder tasks benefit more.
COMPASS dominates the (time, compute) grid on 7/8 hardest tasks; online fine-tuning wins on con-10x10x10a with large budgets.
SGBS leads in the low-budget regime (< 100 s, < 10 parallel attempts) on some RWARE tasks.
Stochastic sampling outperforms online fine-tuning under tight budgets (gradient instability at small batch sizes).
On the single-agent Craftax benchmark, 30-second stochastic sampling alone yields a 37% boost (Appendix C).

Limitations¶

Focus is on multi-agent RL; single-agent evidence is preliminary (one benchmark, one strategy).
Inference strategies are evaluated in isolation — no combinations (e.g., COMPASS + SGBS) are studied.
Out-of-distribution generalization of inference strategies is not analyzed.
Sim-to-reality gap: score functions are assumed accessible or well-approximated during inference, which may not hold for all real-world domains (e.g., wet-lab protein design).
COMPASS requires an additional training phase to construct the latent space, adding pre-deployment cost.
Evaluation hardware is fixed (A100 80 GB); results may not directly transfer to edge/embedded deployment.

Relevance to Agentic AI / LLM Agents¶

This paper is a direct empirical counterpart to inference-time scaling work in LLMs (o1/o3 style "thinking more"), extending test-time compute scaling to multi-agent RL on combinatorial tasks — the class of problems closest to real agentic workflows (logistics, scheduling, routing). The Dec-POMDP formulation and the COMPASS latent-space search strategy are directly applicable to multi-agent LLM systems where diverse strategy rollouts can be scored and selected. The finding that gains are superlinear in task complexity is particularly important for agentic AI researchers: as tasks become more realistic and hard, inference-time search becomes not optional but necessary, echoing the shift from zero-shot to reasoning-augmented LLM agents.