Source-linked AI summary
Exploring Model-based Planning with Policy Networks
Tingwu Wang, Jimmy Ba
TL;DR
Existing MBRL planners rely on random action-space search, which is inefficient in complex environments. POPLIN combines policy networks with online neural-network optimization, achieving state-of-the-art MuJoCo performance and about 3x greater sample efficiency than PETS, TD3, and SAC. Distilled policy networks can also provide effective control without expansive online planning in some environments.
Problem
Existing MBRL planning methods randomly search action sequences, limiting efficiency in complex, high-dimensional environments.
Method
POPLIN combines policy networks with model-based online planning by optimizing action sequences or policy-network parameters.
Results
POPLIN achieves state-of-the-art MuJoCo performance and is about 3x more sample efficient than PETS, TD3, and SAC.
Takeaways & Limitations
Distilled POPLIN policy networks can achieve high performance without expansive online planning in environments such as Cheetah.
Takeaways & Limitations
Action-space policy planning provides little improvement in complex environments because constrained CEM search may fail to converge effectively.
Abstract
from arXiv · showhide
Model-based reinforcement learning (MBRL) with model-predictive control or online planning has shown great potential for locomotion control tasks in terms of both sample efficiency and asymptotic performance. Despite their initial successes, the existing planning methods search from candidate sequences randomly generated in the action space, which is inefficient in complex high-dimensional environments. In this paper, we propose a novel MBRL algorithm, model-based policy planning (POPLIN), that combines policy networks with online planning. More specifically, we formulate action planning at each time-step as an optimization problem using neural networks. We experiment with both optimization w.r.t. the action sequences initialized from the policy network, and also online optimization directly w.r.t. the parameters of the policy network. We show that POPLIN obtains state-of-the-art performance in the MuJoCo benchmarking environments, being about 3x more sample efficient than the state-of-the-art algorithms, such as PETS, TD3 and SAC. To explain the effectiveness of our algorithm, we show that the optimization surface in parameter space is smoother than in action space. Further more, we found the distilled policy network can be effectively applied without the expansive model predictive control during test time for some environments such as Cheetah. Code is released in https://github.com/WilsonWangTHU/POPLIN.
1 Introduction
Existing MBRL planning methods use random action-space search that becomes inefficient in complex, high-dimensional tasks. POPLIN combines policy networks with online planning and reports stronger search efficiency, state-of-the-art benchmark performance, and fast policy-only control in some environments.
- MBRL can improve sample efficiency through learned dynamics and online planning, but scaling and accumulated modeling errors limit practical applications.
- Random shooting and PETS generate candidate action sequences, yet PETS is less effective in higher-dimensional environments.
- Figure 1 compares PETS and POPLIN across planning iterations, showing PETS near initialization while POPLIN searches more efficiently.
- POPLIN uses policy networks to generate proposals for model-predictive control in high-dimensional locomotion tasks with unknown dynamics.
- POPLIN formulates planning as neural-network optimization and reports about 3x greater sample efficiency than PETS, TD3, and SAC.
- POPLIN’s policy-only distilled controller achieves high performance in environments such as Cheetah without expansive online planning.
2 Related Work
Prior MBRL work spans model-based controller optimization, guided policy search, and stochastic value-gradient methods. These approaches include Gaussian-process dynamics, local controllers distilled into policies, and policy optimization through differentiable models.
- Dyna alternates real-environment sampling with controller optimization on a learned environment model.
- PILCO models dynamics with Gaussian processes and directly optimizes surrogate expected reward, but suffers from the curse of dimensionality.
- Guided policy search uses iLQG as a local controller and distills its knowledge into a policy neural network.
- SVG optimizes a stochastic policy network by back-propagating stochastic value gradients through the model.
3 Background
The paper frames reinforcement learning as an infinite-horizon discounted Markov decision process and describes random-shooting planning over learned dynamics. Candidate action sequences are evaluated by predicted rewards.
- The reinforcement-learning task is modeled as an infinite-horizon discounted Markov decision process with state, action, reward, and transition functions.
- Random shooting learns an ensemble dynamics model from previously collected transitions and generates K candidate action sequences over horizon τ.
- The planner selects the candidate action sequence with the best expected reward under predicted dynamics.
4 Model-Based Policy Planning
POPLIN formulates model-based planning as neural-network optimization, using policy networks to initialize action-space search or to guide parameter-space search. It also supports MPC execution, direct policy control, and policy distillation from planned trajectories.
- Planning formulation: POPLIN defines expected planning rewards over policy-generated action sequences and predicted trajectories from ensemble dynamics models.The trajectories are estimated with particles initialized from the current state.
- Action-space planning: POPLIN-A initializes action sequences with a policy network, then uses Gaussian noise and CEM to refine candidate distributions.POPLIN-A-Init uses the policy only for initialization, while POPLIN-A-Replan repeatedly replans along the changed trajectory.
- Action-space planning: POPLIN-A-Replan may converge faster with perfect models and policies but increases the risk of divergent behaviors.The method replans the controller according to the trajectory induced by the current noise distribution.
- Parameter-space motivation: Action-space planning offers little improvement in complex environments because its CEM optimization is non-convex and can average distinct alternatives into a failing action.The paper illustrates this with an obstacle requiring a robot arm to choose either left or right rather than straight.
- Parameter-space planning: POPLIN-P adds noise to policy parameters instead of actions, using either consistent noise across the horizon or separate noise at each time-step.POPLIN-P-Uni reduces the flattened noise vector from (τ + 1)|θ| to |θ|; POPLIN-P-Sep uses separate noise per time-step.
- Control and distillation: Because online MPC is time-consuming, POPLIN also distills planned trajectories and evaluates both MPC control and direct policy control.The paper considers behavior cloning, GAN training, and parameter averaging as distillation schemes.
5 Experiments
The experiments evaluate POPLIN variants against established reinforcement-learning algorithms across MuJoCo environments, including benchmark performance, policy control, search behavior, and hyper-parameter sensitivity. POPLIN performs strongly in many environments, while policy control and some complex tasks remain limited.
- 5.1 MuJoCo Benchmarking Performance: POPLIN is compared with model-based and model-free algorithms across 12 MuJoCo environments using benchmark performance curves and summarized results.The evaluation includes PETS, GPS, RS, MBMF, TD3, METRPO, PPO, TRPO, and SAC.
- 5.1 MuJoCo Benchmarking Performance: POPLIN achieves state-of-the-art performance in almost all environments, solving most tasks within 50,000 or 200,000 time-steps.Performance is averaged across four random seeds, with the reported training budget varying by task difficulty.
- 5.1 MuJoCo Benchmarking Performance: POPLIN-A performs best in simpler environments, whereas POPLIN-P is more consistent and significantly better in complex environments such as Ant and Cheetah.POPLIN-A has no obvious gain over PETS on several complex environments, while POPLIN-P cannot efficiently solve Walker and Humanoid.
- 5.2 Policy Control Performance: Policy control generally lags behind MPC control, although POPLIN-P-BC and POPLIN-P-GAN efficiently distill planned trajectories in several tested environments.Policy control is almost random for Ant and Walker2D, while the relative performance of BC and GAN depends on the environment.
- 5.3 Search Effectiveness and Reward Surface: POPLIN-P improves with larger candidate populations, unlike PETS and POPLIN-A, whose action-space search performance does not increase with population size.The comparison studies PETS, POPLIN-A, and POPLIN-P across different candidate population sizes.
- 5.3 Search Effectiveness and Reward Surface: The action-space reward surface is jagged with many local extrema, while POPLIN searches from low-reward regions toward higher-reward regions and deeper POPLIN-P networks produce more multimodal action distributions.The reward surface is visualized with PCA-projected action sequences, and action distributions are also compared across policy-network depths.
- 5.4 Ablation Study: Performance varies across algorithm variants and hyper-parameters, including training scheme, parameter sharing, initial noise variance, and GAN entropy penalty.POPLIN-P-Sep outperforms POPLIN-P-Uni, larger initial variance benefits POPLIN-P-Avg, and POPLIN-P-GAN is highly sensitive to entropy penalty.
6 Conclusions
The paper combines policy networks with model-based planning through POPLIN and studies distillation for faster testing. It reports state-of-the-art performance on MuJoCo benchmarks and frames online planning as neural-network optimization.
- POPLIN combines policy networks with model-based planning for more efficient control.
- The study evaluates different distillation schemes to provide fast controllers during testing.
- POPLIN obtains state-of-the-art performance on MuJoCo benchmarking environments.
A.1 Algorithm Diagrams
The algorithm diagrams describe two POPLIN-A variants that use cross-entropy planning over action-sequence noise and then update dynamics and policy networks. They differ in whether candidate actions are initialized from a reference sequence or the current policy.
- POPLIN-A-Init: POPLIN-A-Init initializes a reference action sequence before adding and optimizing action-sequence noise.
- CEM Planning: Both variants evaluate predicted trajectories, fit the elite candidates’ noise distribution, and execute the first action from the best sequence.
- Training updates: After interaction, both algorithms update the dynamics model and policy using the collected dataset.
- POPLIN-A-Replan: POPLIN-A-Replan generates candidate actions by perturbing the policy output at each predicted time-step.
A.2 Bench-marking Environments
The evaluation expands beyond PETS’s original four environments to additional OpenAI Gym benchmarks and reports performance at specified time horizons. Figure 8 compares POPLIN variants with several model-free and model-based baselines across 12 environments.
- Environment coverage: The experiments add nine environments beyond the four evaluated in the original PETS paper.
- MuJoCo benchmarks: Table 2 records algorithm performance at 200,000 time-steps on OpenAI Gym MuJoCo environments.
- Classic control: Table 3 records algorithm performance at 50,000 time-steps on classic-control environments.
- Task naming: The PETS Cheetah task differs from OpenAI Gym’s HalfCheetah-v1, which the paper names Cheetah-v0.
- Overall comparison: Figure 8 compares POPLIN-P and POPLIN-A with TD3, SAC, PPO, METRPO, PETS, and RS across 12 benchmarking environments.
A.2.1 Fixing the Swimmer Environments
The updated Swimmer environment became unsolvable for almost all algorithms because of a performance local minimum associated with the neck-mounted velocity sensor. The paper distinguishes this issue from the original Swimmer benchmark and provides a fix.
- After a Gym update, the Swimmer environment became unsolvable for almost all algorithms.
- Most algorithms became trapped near a reward of 130, whereas the original solving threshold is around 340.
- The neck-mounted velocity sensor makes Swimmer highly prone to this performance local minimum.
A.3 Full Results of Bench-marking Performance
POPLIN consistently achieves the best performance across almost all benchmark environments, with results reported alongside algorithm-specific time-step usage and extensive variant analyses. The studies examine hyper-parameter sensitivity, random-seed variation, population size, policy distillation, and reward surfaces in action and parameter spaces.
- Benchmarking performance: POPLIN has the best performance across almost all benchmark environments, with final performance and time-step usage reported for all algorithms.The experiments use probabilistic ensembles with expectation propagation for PETS and POPLIN; alternative combinations can be prohibitively expensive.
- Variant performance: POPLIN-P-BC and POPLIN-P-GAN are comparable, while BC performs better on Cheetah and GAN performs better on Cheetah-v0.These environments are described as essentially the same environment with different naming.
- Population size: POPLIN-P searches efficiently even with zero-fixed parameters, but this setting sacrifices policy distillation and therefore cannot operate without planning.The population-size study links the observed efficiency, especially for POPLIN-P-Avg, to optimization efficiency.
- Reward surfaces: POPLIN-A variants search better than PETS, POPLIN-A-Replan searches more efficiently than POPLIN-A-Init, and deeper POPLIN-P networks improve search efficiency.The reward surfaces are examined in action space for PETS and POPLIN-A and in parameter space for POPLIN-P; POPLIN-A remains limited by dynamics-model errors.
- Reward surfaces: The visualizations compare expected-cost surfaces and planned trajectories across CEM updates for PETS and POPLIN-P networks with zero or one hidden layer.The figures include action-space surfaces and trajectory updates across iterations.