Research map / Reinforcement learning / Offline RL
SUMO: Search-Based Uncertainty Estimation for Model-Based Offline Reinforcement Learning
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes SUMO, a search-based uncertainty estimator for model-based offline RL that uses a KNN-based particle estimate of cross entropy between model and dataset dynamics, plugged into MOPO and AMOReL.
Problem
Model ensemble uncertainty estimates used in model-based offline RL to penalize rewards or truncate synthetic trajectories are unreliable and tied to training, so a better uncertainty estimator is needed.
Approach
Measure uncertainty of synthetic transitions as the cross entropy between model dynamics and dataset dynamics, estimated with a particle-based entropy estimator that reduces to a k-nearest-neighbor distance search (via FAISS) over concatenated (s, a, s') vectors, requiring no extra neural network training.
Key results
- MOPO+SUMO averages 69.3 normalized score vs 59.7 for vanilla MOPO across 15 D4RL MuJoCo datasets, outperforming on 11 of 15.
- AMOReL+SUMO averages 80.3 vs 72.2 for AMOReL; MOBILE+SUMO averages 88.2 vs 85.6.
- SUMO achieves best Spearman or Pearson correlation with model error in 12 of 15 datasets compared with max aleatoric, max pairwise diff and LOO KL.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Replace MOPO uncertainty estimation with SUMO | MOPO · D4RL MuJoCo v2 (15-task average) | normalized score: 59.7 -> 69.3 | Helped |
| Replace AMOReL uncertainty estimation with SUMO | AMOReL · D4RL MuJoCo v2 (15-task average) | normalized score: 72.2 -> 80.3 | Helped |
| Replace MOReL uncertainty estimation with SUMO | MOReL · D4RL MuJoCo v2 (15-task average) | normalized score: 70.7 -> 75.1 | Helped |
| Add SUMO uncertainty estimation | MOBILE · D4RL MuJoCo v2 (15-task average) | normalized score: 85.6 -> 88.2 | Helped |
| Replace MOPO uncertainty estimation with SUMO | MOPO · D4RL Antmaze v0 (6-task average) | success rate: 0.0 -> 9.2 | Helped |
| Replace AMOReL uncertainty estimation with SUMO | AMOReL · D4RL Antmaze v0 (6-task average) | success rate: 0.0 -> 3.3 | Helped |
| Include action and next state in the search vector (s⊕a⊕s′ instead of s⊕s′) (ablation) | SUMO · D4RL MuJoCo (6 selected datasets) | Spearman rank correlation (ρ): 0.56 -> 0.83 | Helped |
| Include action and next state in the search vector (s⊕a⊕s′ instead of s⊕a) (ablation) | SUMO · D4RL MuJoCo (6 selected datasets) | Spearman rank correlation (ρ): 0.63 -> 0.83 | Helped |
| Replace Euclidean distance with Manhattan distance (ablation) | MOPO+SUMO · D4RL MuJoCo (6 selected datasets) | normalized score (average): 75.8 -> 75.5 | No clear effect |
| Replace Euclidean distance with Cosine similarity (ablation) | MOPO+SUMO · D4RL MuJoCo (6 selected datasets) | normalized score (average): 75.8 -> 74.3 | No clear effect |
| Set KNN neighbor count to 5 instead of 1 (ablation) | SUMO · D4RL MuJoCo (6 selected datasets) | Spearman rank correlation (ρ): 0.83 -> 0.82 | No clear effect |
| Set KNN neighbor count to 10 instead of 1 (ablation) | SUMO · D4RL MuJoCo (6 selected datasets) | Spearman rank correlation (ρ): 0.83 -> 0.81 | No clear effect |
Surprising findings
- Using MOPO+SUMO's KNN uncertainty with k=1 gives comparable performance to larger k, indicating robustness to k.
- Including the next state in the search vector substantially improves correlation over (s,a) only; the paper also notes that Manhattan and cosine distances perform similarly to Euclidean.
Limitations
- Heavy computational overhead of KNN search as data dimension and dataset size grow, acknowledged in conclusion.
- Antmaze results are near zero for most tasks, so gains there are limited.
- Theoretical bounds rely on assumptions; the cross-entropy estimate is an approximation via particle estimator.
Experimental setup
Datasets: D4RL halfcheetah D4RL hopper D4RL walker2d D4RL antmaze
Benchmarks: D4RL MuJoCo D4RL Antmaze
Metrics: normalized score success rate Spearman rank correlation Pearson correlation
Baselines: MOPO MOReL MOBILE AMOReL COMBO RAMBO Max Aleatoric Max Pairwise Diff LOO KL Divergence EDAC DARL
Builds on: MOPO MOReL SAC FAISS
Closest papers by meaning
- OCEAN-MBRL: Offline Conservative Exploration for Model-Based Offline Reinforcement Learning (AAAI 2024 · similarity 0.91)
- Dynamic Uncertainty Estimation for Offline Reinforcement Learning (AAAI 2025 · similarity 0.91)
- A Unified Principle of Pessimism for Offline Reinforcement Learning under Model Mismatch (NeurIPS 2024 · similarity 0.90)
- Selective Uncertainty Propagation in Offline RL (AAAI 2025 · similarity 0.90)
- The Edge-of-Reach Problem in Offline Model-Based Reinforcement Learning (NeurIPS 2024 · similarity 0.89)
- Model Selection for Off-policy Evaluation: New Algorithms and Experimental Protocol (NeurIPS 2025 · similarity 0.89)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.