The work, published in Machine Learning, addresses a problem hidden inside many offline reinforcement-learning systems. These systems first learn an approximation of the environment’s dynamics from historical data. The approximation predicts what will happen after an agent takes an action: how the state will change, what reward may be received and how uncertainty grows away from familiar situations. The agent then uses this learned world to generate imaginary experience and train a policy. If the chosen model is subtly inaccurate, however, the policy may learn to exploit errors in the simulation rather than useful behavior in the real environment. A model that achieves the lowest prediction error on a held-out dataset is not necessarily the model that produces the best real-world policy, particularly when the policy visits states and actions missing from the original data.

The authors, Yu-Wei Yang, Wei Hung, Yun-Ming Chan, Xi Liu and Ping-Chun Hsieh, focus on this mismatch between conventional model evaluation and actual policy performance. Standard validation treats dynamics learning much like supervised learning: a model is rewarded for accurately predicting recorded transitions. But an offline dataset covers only a limited slice of the possible state-action space. A policy trained on the model may travel beyond that slice, where validation accuracy offers little reassurance. Another option is off-policy evaluation, which estimates how well the policy would perform without directly deploying it. Yet off-policy evaluation also becomes unreliable when the dataset does not sufficiently cover the policy’s behavior. In the researchers’ motivating experiments, both validation and off-policy evaluation could select models whose resulting returns were far below those of the best available candidate.

BOMS changes the procedure by treating model selection as a black-box optimization problem. The candidate models are the input points, and the unknown objective is the true return of the policy trained with each model. Evaluating that objective requires running the policy in the actual environment, making every evaluation expensive and potentially risky. Bayesian optimization is designed for precisely this situation. It builds a probabilistic surrogate of an unknown function from a small number of observations, then chooses new points by balancing exploitation—testing models believed to be strong—with exploration—testing models about which the surrogate remains uncertain. BOMS uses a Gaussian process and the GP-UCB acquisition rule to decide which dynamics model should be evaluated next. After each evaluation, the observed return updates the surrogate, refining the search.

The crucial technical challenge is defining how similar two dynamics models are. Neural-network models cannot be compared reliably simply by measuring the distance between their weights. Two networks with very different parameters may make similar predictions, while small parameter changes can sometimes produce significant behavioral differences. BOMS therefore constructs a model-induced distance based on predictions in regions relevant to the current candidate policy. Given a batch of states from the offline dataset, the method samples actions from the policy trained using the currently selected model. It then compares the next-state predictions and reward predictions made by that model and every other candidate. The distance combines the expected difference between predicted next states with a weighted difference between predicted rewards. In the reported experiments, equal weighting of these two terms was generally effective.

This design is motivated by a theoretical performance-gap result. Under a Lipschitz-continuity assumption on the relevant value functions, the researchers bound the difference between the true-environment returns of policies learned from two uncertainty-penalized models. The bound contains factors connected to the behavior policy and its modeling error, as well as a term measuring the discrepancy between the two models’ predicted next states and rewards. The behavior-policy terms do not change when comparing candidate models, so the prediction discrepancy becomes the practical signal for constructing a similarity measure. In other words, models that make similar one-step predictions under policy-relevant conditions are treated as nearby points in the Bayesian-optimization search space. The resulting radial-basis-function kernel allows BOMS to infer the likely performance of unevaluated models.

The experiments tested the framework across a broad collection of control problems. The researchers trained candidate models following the procedure used by MOPO, a prominent offline model-based reinforcement-learning algorithm, and assembled 150 models for each task. The evaluation included MuJoCo locomotion environments—walker2d, hopper and halfcheetah—using medium, medium-replay and medium-expert datasets. It also included the Adroit pen-manipulation task and the Meta-World door-opening task. At each Bayesian-optimization step, the policy associated with a selected model was evaluated through a small number of real-environment trajectories. For the locomotion tasks, the evaluation averaged five trajectories; for the Adroit pen task, it used 20 shorter trajectories. The entire selection process was repeated across multiple trials to account for randomness in initialization, policy training and Monte Carlo return estimates.

The central result was strikingly economical: BOMS often delivered substantial gains after only a handful of online evaluations. The study reports that roughly five selection iterations, representing approximately 1% to 2.5% of the offline training data in the tested settings, were enough to outperform validation-based selection and fitted Q-evaluation in nearly all tasks. Across the full process, the researchers used 20 Bayesian-optimization iterations, corresponding to only a few dozen evaluation trajectories in many environments. BOMS also outperformed random selection, suggesting that the improvement was not simply the result of eventually trying several models. Its surrogate model was able to exploit structure among the candidates, moving toward neighborhoods where high-performing models were concentrated. The approach also improved RAMBO, another offline model-based reinforcement-learning method, indicating that the selection strategy is not limited to MOPO.

Additional tests revealed why the particular distance measure matters. Comparing neural-network weights performed poorly, as did distances based on unrelated policies or random exploratory actions. Distances generated using the policy trained by the currently selected model were generally more useful because they focused comparisons on states and actions that matter for the decision being made. Surprisingly, extending the distance calculation to five- or 20-step model rollouts did not improve performance. Longer imagined rollouts accumulate prediction errors, a familiar problem in model-based reinforcement learning, and can make two models appear different because of compounded inaccuracies rather than meaningful one-step disagreement. The one-step distance offered a more stable signal, although the authors acknowledge that it may become less reliable under severe distribution shift or when learned models are poor in the regions visited by candidate policies.

The findings do not turn offline reinforcement learning into a risk-free technology. BOMS requires online interaction, so it is not applicable to strictly offline settings in which no environmental evaluation is permitted. Nor does the method itself guarantee safe execution of the policies it queries. In safety-critical applications, the authors suggest combining it with action shielding, constrained policy execution, early termination rules, conservative fallback policies or other safety layers. The candidate pool may also become expensive to search when it contains extremely large numbers of models, although the Gaussian-process calculations scale linearly with the number of candidates during posterior evaluation and can be parallelized. The researchers propose pre-screening, clustering and sparse Gaussian-process approximations as possible solutions. Even with these limitations, BOMS offers a compelling shift in perspective: instead of asking which model predicts the logged data most accurately, offline reinforcement learning can ask which model produces the most trustworthy behavior—and answer that question with a remarkably small amount of real-world evidence.

Subject of Research: Active model selection for offline model-based reinforcement learning using Bayesian optimization

Article Title: Enhancing Offline Model-Based RL via Active Model Selection: A Bayesian Optimization Perspective

Article References: Yang, Y.-W., Hung, W., Chan, Y.-M., Liu, X. et al. “Enhancing Offline Model-Based RL via Active Model Selection: A Bayesian Optimization Perspective.” Machine Learning, 115, Article 179 (2026). Key references include Yu et al. (2020), “MOPO: Model-based offline policy optimization”; Rigter et al. (2022), “RAMBO-RL: Robust adversarial model-based offline reinforcement learning”; and Konyushova et al. (2021), “Active offline policy selection.”

Image Credits: AI Generated

DOI: 10.1007/s10994-026-07119-6

Keywords: Offline reinforcement learning, Bayesian optimization, dynamics model selection, model-induced kernel, Gaussian process, active model selection, model-based reinforcement learning, distribution shift, uncertainty estimation, policy evaluation

Traditional model-based offline RL approaches typically employ simplistic predictive models that map current states and actions directly to predicted next states. While seemingly straightforward, such techniques are susceptible to capturing spurious correlations that arise due to the inherent biases in the offline datasets, which are often influenced by the sampling policies that generated the data. These misleading correlations can degrade generalization capabilities, producing policies that perform poorly when confronted with previously unseen situations. Recognizing these limitations, the research team argues for a paradigm shift that emphasizes causal inference as a more principled foundation for model learning within offline RL.

Central to their proposition is the notion that environment models should encapsulate the underlying causal influences among state variables and actions. By explicitly uncovering causal dependencies, these models can potentially disentangle genuine mechanisms driving state transitions from confounding statistical artifacts, thereby facilitating the development of policies that generalize robustly beyond the offline data distribution. To address this, the team introduces FOCUS, an acronym for offline model-based reinforcement learning with causal structured world models, which integrates causal discovery with model-based RL algorithms to exploit the causal structure for enhanced policy learning.

FOCUS begins by deriving the causal relationship matrix from given offline data through kernel-based conditional independence testing (KCI test), a nonparametric method that does not assume linearity or specific distributional forms and works effectively with continuous variables. This step aims to identify the most plausible causal connections between state features by analyzing conditional independencies, a key component in causal inference frameworks. Subsequently, FOCUS determines the causal structure by selecting an appropriate threshold on the resulting p-values, thereby constructing a causal graph that encodes the directional dependencies foundational to the environment’s dynamics.

One notable innovation of the FOCUS methodology lies in its exploitation of the temporal nature of reinforcement learning data. By leveraging the fundamental principle that causes precede effects in time, the researchers incorporate a temporal constraint into the PC algorithm, a popular causal discovery method. This constraint, which enforces that future states cannot influence past states, drastically reduces the computational burden by narrowing down the scope of hypothesis testing that the algorithm needs to consider. This is particularly critical given the typically large number of conditional independence tests required in causal discovery, which would otherwise be computationally prohibitive in high-dimensional scenarios.

After unraveling the causal structure, FOCUS merges this insight with a neural network-based environment model, enabling the learned dynamics to be guided by causal principles. This integration facilitates an offline model-based reinforcement learning scheme that trains policies grounded in a causally consistent world model. The research team provides rigorous theoretical evidence demonstrating that such causal environment models yield tighter generalization error bounds compared to plain predictive models, underscoring the statistical advantages of embedding causality into RL frameworks.

Empirical evaluations showcased in the study reveal that FOCUS substantially outperforms baseline offline model-based RL methods and existing causal MBRL algorithms across various benchmark tasks. These findings not only validate the theoretical predictions but also highlight the practical impact of causal discovery in improving policy learning from static datasets. By emphasizing causal inference, FOCUS mitigates the risk of overfitting to spurious correlations and promotes policies with broader generalizability, a critical factor for real-world applications where data is collected offline and interaction is costly or dangerous.

Moreover, the study underscores broader implications for the field of artificial intelligence by illustrating how causality can be systematically integrated into reinforcement learning to overcome fundamental challenges posed by data biases and confounding factors. As AI systems increasingly enter safety-critical domains, from autonomous driving to healthcare, ensuring that learned policies are causally sound and reliable is paramount. The FOCUS framework represents an important step in this direction, combining statistical rigor with computational efficiency.

The researchers emphasize that while causal discovery is inherently challenging due to the combinatorial explosion of potential hypotheses, cleverly leveraging domain-specific properties such as temporal order can make the problem tractable in practical scenarios. This insight has the potential to influence future developments in causal reinforcement learning, inspiring new algorithms that refine causal structure learning under operational constraints. Additionally, the adoption of kernel-based conditional independence tests broadens the applicability of FOCUS to diverse data types encountered in real-world tasks.

This work was published on April 15, 2025, in the journal Frontiers of Computer Science, co-published by Higher Education Press and Springer Nature. It represents a collaborative effort between experts specialized in causal inference, reinforcement learning, and machine learning theory, contributing substantially to the ongoing dialogue on bridging causality and artificial intelligence. The publication further cements LAMDA’s role as a pioneering research institution advancing foundational AI methodologies.

The study’s findings open intriguing avenues for future research, including extending FOCUS to online RL settings, incorporating richer causal models with latent confounders, and exploring transfer learning scenarios where causal structures discovered in one domain inform policy learning in another. Such endeavors will continue to clarify how humans’ innate causal reasoning abilities can be emulated and leveraged by artificial agents for more robust decision-making.

In conclusion, the introduction of FOCUS marks a significant advancement in offline reinforcement learning by directly addressing the limitations of conventional predictive models through a principled incorporation of causal discovery. By marrying causal inference techniques with neural network–based environment modeling and offline policy optimization, this approach sets new standards for learning reliable, generalizable policies from static datasets, paving the way for more trustworthy and effective AI systems in complex, real-world environments.

Subject of Research: Not applicable

Article Title: Offline model-based reinforcement learning with causal structured world models

News Publication Date: 15-Apr-2025

Web References:
https://doi.org/10.1007/s11704-024-3946-y

Image Credits: Zhengmao ZHU, Honglong TIAN, Xionghui CHEN, Kun ZHANG, Yang YU

Keywords: Computer science