Boosting MCTS with Free Energy Minimization

TL;DR

The paper tackles planning under uncertainty for continuous-control tasks by unifying Monte Carlo Tree Search with Active Inference. It introduces MCTS-CEM, a framework that fits a single Gaussian action distribution at the root with the Cross-Entropy Method and reuses it throughout tree search and rollouts, while incorporating an epistemic (information-gain) term into the objective. Key contributions include a principled blend of extrinsic rewards and intrinsic exploration, a BALD-inspired approximation for epistemic value, and empirical gains over pure CEM and MCTS with random rollouts across several continuous-control benchmarks. The work demonstrates robust planning performance in sparse and high-dimensional tasks, and points to future directions like adaptive intrinsic exploration and deeper horizon planning to further strengthen active-inference-guided search.

Abstract

Active Inference, grounded in the Free Energy Principle, provides a powerful lens for understanding how agents balance exploration and goal-directed behavior in uncertain environments. Here, we propose a new planning framework, that integrates Monte Carlo Tree Search (MCTS) with active inference objectives to systematically reduce epistemic uncertainty while pursuing extrinsic rewards. Our key insight is that MCTS already renowned for its search efficiency can be naturally extended to incorporate free energy minimization by blending expected rewards with information gain. Concretely, the Cross-Entropy Method (CEM) is used to optimize action proposals at the root node, while tree expansions leverage reward modeling alongside intrinsic exploration bonuses. This synergy allows our planner to maintain coherent estimates of value and uncertainty throughout planning, without sacrificing computational tractability. Empirically, we benchmark our planner on a diverse set of continuous control tasks, where it demonstrates performance gains over both standalone CEM and MCTS with random rollouts.

Figures (11)

• Figure 1: MCTS-CEM Diagram. A: Initialize the MCTS tree with the agent's current state $s_0$. B: Fit the root node's action distribution using CEM. Actions are evaluated by minimizing expected free energy ($G_i$), with next states sampled using the current action Gaussian. The epistemic value $EV^i_t$ is computed as the KL divergence, and rewards ($r_t^i$) approximate $\ln P(y_t^i)$. The top-performing actions refine the distribution iteratively. C: Use the fitted action distribution for action sampling during MCTS exploration, balancing exploitation and exploration with UCB-like selection and consistent simulations at the leaves.
• Figure 2: Root Action Distribution Fitting Using CEM: This diagram focuses on subsection 3 of component B in the MCTS-CEM process. Candidate actions $\{a^{(1)}, \ldots, a^{(n_{\text{candidates}})} \}$ are sampled from a Gaussian distribution. Their evaluations, based on the expected free energy objective $G_i$, approximate extrinsic value $ln P(y_t^i) \approx r^i_t$ and epistemic value $EV^i_t$ from KL divergence. The top $k$ candidates refine the distribution, optimizing exploration and exploitation.
• Figure 3: MCTS-CEM Planning (Component C). After fitting the root action distribution, MCTS uses it to drive action sampling at each expansion step and for leaf-node simulations.
• Figure 4: The Pendulum environment, where the agent applies torque to swing the pendulum to its upright position.
• Figure 5: Performance comparison of MCTS-CEM, CEM, and MCTS-Random on the Pendulum environment, showing the cumulative reward over episodes averaged over five trials with different random seeds. Error bars represent the standard deviation across trials. In this well-shaped, deterministic setting, the additional exploration provided by MCTS-CEM results in performance comparable to CEM, indicating limited added benefit beyond a straightforward Cross-Entropy Method approach.