PaperlayerView exact revision on arXiv

Source-linked AI summary

Adaptive Mixing of Policies from Searching and Policies from Learning

Gavin B. Rens

arXiv:2608.15700v1cs.AI

TL;DR

Planning can be computationally expensive, motivating adaptive coordination between learned policies and search. Flexer mixes neural-network and MCTS policies according to policy and model quality, outperforming alternatives on selected harder toy problems.

  • Problem

    Planning in reinforcement learning is computationally expensive, creating a need to reuse planning while adapting its extent to policy and model quality.

  • Method

    Flexer mixes policy-network and MCTS outputs at each step using quality-based weighting and adaptively adjusts planning parameters.

  • Results

    On hard MovingNumbers, only Flexer-BT solves the problem, while performance is comparable on easy environments and AZLike leads on BlocksWorld[4,4].

  • Takeaways & Limitations

    Flexer shows a scalability advantage on complex, sequential planning tasks within the tested symbolic benchmarks.

  • Takeaways & Limitations

    Flexer is more complex than AlphaZero/AZLike, with interacting components, additional hyperparameters, and potential failure modes.

Abstract

from arXiv · show

Background: Distillation of training targets generated thru search/planning has proven useful in reinforcement learning, but search can take exceedingly long. Objectives: Rather than perform search to the same depth every time (typically at a fixed period of steps), reduce the search depth proportionally to the quality of the policy network priors. Methods: We describe Flexer, an architecture that, for each step, mixes the policy from a neural network and the policy from Monte Carlo tree search. The mixing factor favors the MCTS policy as the policy imitation error of the network and the environment models' variance increases. Results: Flexer outperforms a version of AlphaZero (and DQN and ADP) for some experiments on three toy symbolic problems.

1 Introduction

Flexer is a model-based reinforcement-learning architecture that adaptively mixes neural-network and MCTS policies, reducing planning when policy or model quality is high. Evaluations on three small problems show mixed results but indicate potential to reduce wasted planning computation.

  • Architecture: Flexer adds policy trade-off and planning-budget control to AZLike, an AlphaZero-based architecture for single-agent stochastic environments with learned models.Unlike AlphaZero and MuZero, Flexer maintains a global value function that mediates information between the policy network and planner.
  • Results: Across three small problems, Flexer results are mixed but suggest that adaptive planning can reduce computation wasted on planning.The experiments also include Double Deep Q-learning Network (DQN) and Adaptive Dynamic Programming (ADP) as additional perspectives on the model-based approach.
  • Motivation and contributions: Flexer adaptively mixes MCTS and policy-network policies using a factor inversely proportional to policy-network and environment-model quality.The architecture also adjusts MCTS temperature and exploration according to the mixing factor and policy quality.
  • Evaluation: The study evaluates value-function intermediation, adaptive policy mixing, interior-node value boosting, and their dependence on environment complexity.These effects are examined across three toy environments and compared with comparable AlphaZero and standard MuZero versions.

2 Basic Concepts and Related Works

This section introduces MDPs and MCTS, then situates Flexer relative to AlphaZero, AZLike, MuZero, and related model-based or model-free approaches. Flexer’s context differs from AlphaZero because it learns stochastic environment models for single-player problems without access to real dynamics.

  • Markov decision processes: An MDP comprises states S, actions A, transition function T, reward function R, and initial state s0.Transitions specify probabilities between states, while rewards reflect the utility of those transitions.
  • Monte Carlo Tree Search: MCTS searches forward from s0 by iteratively selecting, expanding, rolling out, and backpropagating values through a tree.Each branch stores an estimated state-action value Q̂(s, a), updated during backpropagation; standard selection often uses UCT.
  • AlphaZero foundations: AlphaZero trains a policy-value network from MCTS data, selecting actions from root-child visit counts while guiding search with learned policy priors.Its planner assumes access to known transition and reward models.
  • AZLike baseline: AZLike is the paper’s AlphaZero-based baseline for single-player problems with stochastic transitions and rewards learned without access to real environment dynamics.The architecture is presented conceptually in Figure 2.
  • Related model-based methods: MuZero similarly learns environment models for planning without a simulator, using representation, dynamics, and prediction networks and latent-space MCTS.Unlike AlphaZero, MuZero normalizes unbounded search-tree Q-values into [0, 1] before applying modified PUCT.

3 The Proposed Architecture: Flexer · 3.1 Mixing Logic

Flexer comprises four variants that differ in planner rollout and environment-model representation. At each decision, it mixes MCTS and policy-network action distributions using a factor determined by policy and environment-model quality.

  • 3 The Proposed Architecture: Flexer: Flexer has four variants combining rollout versus direct value-function bootstrapping with tabular versus neural-network environment models.The rollout variants use MCTS rollout, while the others bootstrap newly generated node values from the value function in AlphaZero style.
  • 3.1 Mixing Logic: At each step, Flexer selects the action with the highest value from a mixture of the MCTS and policy-network action distributions.The mixture is represented as π(s), combining πMCTS(s) and πPN(s).
  • 3.1 Mixing Logic: The mixing factor µ ∈ [0, 1] controls the ratio between the MCTS and policy-network distributions.Flexer determines µ from the current qualities of the policy network and environment model.
  • 3.1 Mixing Logic: µ should favor the policy network as κPN improves, and favor MCTS when κPN is low but κEM is high.When both sources are unreliable, the agent should take a random action.
  • 3.1 Mixing Logic: Policy-network quality is defined as κPN = 1 − ψPN, where ψPN is normalized policy imitation error.ψPN measures policy-network inference quality.
  • 3.1 Mixing Logic: For tabular environment models, κEM is based on the smaller of 1 − Tvar and 1 − Rvar, reflecting transition-probability and reward-history variance.The MCTS policy is relied on in proportion to the least reliable of these two variance-derived values.
  • 3.1 Mixing Logic: Flexer combines policy and environment-model quality as x = (2 − κPN − κEM)/2, then reshapes it into rand act = (e^kx − 1)/(e^k − 1) with k = 10.x ranges from 1 at the worst combined quality to 0 at the best, while the reshaped function maps [0, 1] to [0, 1].

3.2 Definitions of Mixing-Factor Factors

Flexer derives its mixing factors from two smoothed quality signals: policy imitation error and local environment-model uncertainty. These signals use KL divergence for policy distributions and prediction variance or accuracy for learned transition and reward models.

  • Policy quality: Policy imitation error (PIE) is the batch-average KL divergence between target action distributions and the current policy network distributions.The policy distribution is defined as pi = softmax(network(si)).
  • Policy quality: PIE is normalized by log |A| and tracked with an exponential moving average ψP N, using a smoothing factor such as α = 0.05 to 0.2.The persistent estimate represents the network’s long-term distance from its training targets.
  • Environment-model quality: Flexer estimates local transition and reward uncertainty by filtering historical buffers to states near the current state and computing sample variances over observed outcomes.Local states are identified by breadth-first expansion from the current state using the learned transition model.
  • Environment-model quality: Exponentially smoothed transition and reward variance trackers, Tvar and Rvar, provide model-reliability information to the mixing function µ(·).The variance updates use step-size α, with an example value of 0.05.
  • Metric choice: Environment-model quality is measured with tolerance-based next-state prediction accuracy, whereas policy quality uses KL divergence against the MCTS visit-count distribution.The state-prediction metric uses tolerance δ = 0.05, reflecting the different output types of the two networks.

3.3 The State Value Function

Flexer’s state value function is updated online from agent experience and optionally from MCTS state values. It supports leaf evaluation, policy-network training data, and MCTS rollout decisions while using an incrementally learned empirical model.

  • Value-function role: The state value function V is updated at every agent step from experience and optionally from state values computed in the MCTS tree.These updates are central to the Flexer architecture.
  • Value-function role: V is used at MCTS leaf nodes, to generate policy-network training data, and in the MCTS rollout policy.
  • Online value update: After each observed transition, Flexer updates its empirical transition and reward models before performing a single-step incremental value-iteration update for V.The update approximates the Bellman optimality operator using the learned models.
  • Model fallback: If an action has never been tried from state s, the algorithm uses a prior that keeps the agent in s and applies the constant step cost.An example given is rstep = −1.
  • Online value update: The incremental update is mathematically equivalent to standard value iteration with learning rate α and supports smooth online convergence while the model is populated.The resulting value is then added to the value buffer.

3.4 Value-function Boosting from Interior MCTS Nodes

Flexer can boost a global tabular value function using targets from all visited internal MCTS nodes, with either visit-weighted or maximum child Q-values. Updates use a depth-discounted learning rate, extending prior approaches that update only roots or selected paths.

  • Method: Flexer recurses through all visited internal MCTS nodes and updates a global value function using either maximum-Q or visit-weighted-Q targets.The two node-value variants are based on maximum Q-values leading to children or visit-weighted Q-values of children.
  • Method: Updates apply the learning rate α / (depth+1), so deeper nodes receive smaller update weights.At each node, the value target is computed from local Q-statistics before applying the depth-discounted learning rate.
  • Related work: A0GB likewise uses internal-tree Q-values as value targets, but specifically selects the leaf of a greedy path rather than all visited internal nodes.A0GB was reported to find optimal policies in tabular domains where original AlphaZero fails and to train faster on Connect-Four.
  • Novelty: The authors found no published work combining all internal-node updates with a depth-discounted learning rate.This distinction is presented as the method’s novelty relative to the surveyed prior work.

3.5 Policy Network Training

The policy network is trained by regressing its outputs toward greedy, value-function-based action targets generated from learned environment models. Training uses buffered state-target pairs, periodic minibatch updates with entropy regularization, and an adaptive number of epochs based on policy quality.

  • Training targets: The policy network minimizes MSE between its raw logits and hard one-hot greedy targets induced by V(s), learned transition model P̂, and reward model R̂.The resulting (state, target) pairs are stored in the policy-network experience buffer BPN.
  • Optimization: When the experience buffer reaches its minimum size, training samples batches and adds a small entropy regularizer to the MSE objective.Training therefore begins only after sufficient state-target data have been accumulated.
  • Training schedule: Training calls occur periodically, once per episode in this study, and each call runs for a fixed number of epochs.Flexer instead adapts the epoch count to current policy quality ψPN.
  • Training schedule: Flexer uses 3 base epochs plus up to 7 extra epochs, for 3–10 total epochs, with extra training determined by policy imitation error.The implementation sets extra_epochs = max(0, int(7 * PIE)).
  • Target collection: For each tree node with at least one child, the method calculates an action probability distribution and adds the resulting state-target pair to BPN.The stored state is the state represented by the node.

3.6 Flexer’s MCTS

Flexer computes an MCTS policy using PUCT-based tree search with rollout or value bootstrapping, then modulates its planning budget and softmax temperature through the mixing factor µ. As µ increases, MCTS receives more iterations and a higher temperature, while AZLike uses half Flexer’s maximum iteration count.

  • MCTS search procedure: PUCT descends through fully expanded nodes, then expands an untried action uniformly and evaluates the new child by rollout or value bootstrapping.Rollout depth is limited by the configured maximum tree depth.
  • MCTS search procedure: Neural environment-model variants avoid repeated inference by sampling successors from experienced transitions stored alongside the neural representation.If no experienced successor exists, the successor state is set equal to the current state.
  • Mixing-factor modulation: The MCTS policy temperature is τ = max{0.01, 0.2µ}, increasing with µ so greater planning budget preserves less-uniform visit-count information.The method argues that lower µ should retain uncertainty because reduced planning is more likely to produce uniform visit counts.
  • Mixing-factor modulation: The planning budget weights the maximum MCTS iteration count by µ, whereas AZLike always uses half the maximum Flexer iteration count.As µ approaches 1, planning is intended to dominate and MCTS should use the maximum user-defined iterations.

3.7 A Flexer Algorithm

Flexer combines a policy network, a value function, and an MCTS planner, with Algorithm 3 specifying the main procedure. Flexer-RN and Flexer-BN differ by estimating environment-model quality from the neural-network environment model and updating it once per episode.

  • Architecture: Flexer adds a value function that interacts with the policy network and MCTS planner.This distinguishes the architecture from AZLike.
  • Training loop: The training loop uses ϵ-greedy exploration for all methods except ADP, with ϵ decaying from 1 to 0 across episodes.The study notes this exploration choice is likely suboptimal and does not investigate alternatives for Flexer.
  • Algorithm variants: Algorithm 3 presents the Flexer-RT and Flexer-BT versions.The passage identifies Algorithm 3 as containing most of the main contributions.
  • Algorithm variants: Flexer-RN and Flexer-BN compute κ_EM from the neural network representing the environment model rather than transition-and-reward variance.Their κ_EM is computed and maintained in an environment-model training process that runs once per episode.

4 Evaluation

Flexer and AZLike were evaluated across three symbolic environments, using easy and harder variants with success, return, efficiency, and action-time measures. The evaluation compared all variants on easy tasks and selected best-performing variants on harder tasks.

  • Evaluation: All Flexer and AZLike variants were tested on easy environments; best variants were then tested on harder environments.The harder-environment AZLike runs used a third the training epochs and half the number of ...
  • Evaluation: MCTS settings were fixed per environment instance and chosen as small as reasonably possible, with maximum iterations of 50 for Flexer and 25 for AZLike.Maximum tree depth was 15 with rollout and 25 with bootstrapping; rollout depth was at most 10.
  • Evaluation: Performance was measured by successful completions, maximum sustained return, steps per episode, and time per action.Fewer steps indicate earlier success.
  • Evaluation: Episodes allowed at most 200 steps, with averages reported over 30 easy-problem instances and 60 hard-problem instances.These instance counts applied across the experiments described in the evaluation setup.
  • Evaluation: Evaluation covered SimpleGrid, BlocksWorld, and MovingNumbers, each with easier and harder versions.SimpleGrid requires reaching a maze goal, BlocksWorld requires stacking blocks in order, and MovingNumbers requires moving numbers to corresponding locations in sequence.

5 Discussion

Discussion results show that neither Flexer nor AZLike consistently dominates across environments, with performance and runtime varying by variant. Value boosting can substantially improve Flexer in some settings, while ADP remains superior on selected tasks.

  • Cross-environment comparisons: AZLike achieves 92 goals versus 85 for Flexer in one environment, while AZLike-BT runs twice as fast as Flexer.Flexer runs twice as fast as AZLike-RN/RT.
  • Cross-environment comparisons: Flexer-RT achieves 126 goals versus AZLike’s 122 out of 150, with Flexer runtimes clustered around 0.015–0.02 per action.All tabular variants improve in goals achieved, and Flexer-BT/RT/RN have non-decreasing returns and non-increasing steps.
  • Cross-environment comparisons: Performance remains environment-dependent: AZLike-RT achieves 106 goals versus Flexer-BT’s 96, while both AZLikes outperform Flexer variants by approximately 40 versus 80 goals out of 300 elsewhere.Flexer-BT performs particularly poorly in the latter environment, and no clear overall winner emerges.
  • Cross-environment comparisons: Flexer-BT significantly outperforms all three other variants on MovingNumbers[10,3], described as arguably the hardest tested problem.The result contrasts with other environments where AZLike or ADP performs better.
  • Value boosting: Value boosting improves Flexer selectively: MQ clearly boosts Flexer-BT, both MQ and WQ boost Flexer-RT, and one case shows no significant effect.Boosting generally adds little runtime cost, although one experiment reports a temporary approximately 27% penalty at 160 episodes.
  • Comparison with ADP and DQN: ADP achieves 150 goals versus Flexer-BT-MQ’s 123 out of 200 in one task, but Flexer-RT-MQ reaches 132 versus ADP’s 124 in another.ADP is much faster, at approximately 0.004 seconds per action versus Flexer’s approximately 0.0175 at training’s end.
  • Comparison with ADP and DQN: Flexer solves a task that ADP and DQN cannot, finding the goal 19 times with BT-MQ versus 12 times with BT-WQ out of 200 episodes.MQ achieves approximately seven more goals on average and finds the goal in fewer than 140 steps, versus over 160 for BT-WQ.

6 Conclusions

Flexer combines adaptive policy mixing, a tabular value function, and value boosting to improve planning efficiency and learning stability. Its strongest results occur on the hardest MovingNumbers task, while complexity, model-quality estimation, and scalability remain challenges.

  • Architecture: Flexer mixes the policy-network output with the planner root result using a quality-dependent factor that favors planning when policy quality is low and model quality is high.The same quality measures adapt parameters to improve performance and save compute budget.
  • Architecture: Flexer derives training targets from a tabular value function rather than planner root nodes and uses MCTS node values to boost that function.This distinguishes it from AlphaZero and MuZero.
  • Limitations and future work: Flexer’s greater complexity introduces interacting components, additional hyperparameters, and potential failure modes that make analysis more challenging.Scaling also requires addressing neural-network inference cost and the state-space-dependent overhead of computing environment-model variances.
  • Results: Flexer-BT decisively outperforms all variants on the hardest MovingNumbers[10,3] problem, whereas AZLike performs better on BlocksWorld[4,4].Bootstrapped values provide gradient signal without goal completion, while adaptive budgeting limits misleading effects from poor early models.
  • Results: Flexer produces more stable learning curves, bootstrapping is the most impactful design choice, and MQ outperforms WQ on hard MovingNumbers at no extra runtime cost.Value-function boosting helps on SimpleGrid and BlocksWorld, while ADP is competitive on simpler environments and near-optimal on BlocksWorld.
  • Limitations and future work: Future work could investigate entropy- or variance-based model-quality measures and latent-space techniques to achieve true scalability.Latent-space approaches are identified through examples including DVRL, MuZero, Dreamer, and BetaZero.
Loading 2608.15700v1…