Procedural Generalization by Planning with Self-Supervised World Models

Abstract

One of the key promises of model-based reinforcement learning is the ability to generalize using an internal model of the world to make predictions in novel environments and tasks. However, the generalization ability of model-based agents is not well understood because existing work has focused on model-free agents when benchmarking generalization. Here, we explicitly measure the generalization ability of model-based agents in comparison to their model-free counterparts. We focus our analysis on MuZero [60], a powerful model-based agent, and evaluate its performance on both procedural and task generalization. We identify three factors of procedural generalization-planning, self-supervised representation learning, and procedural data diversity-and show that by combining these techniques, we achieve state-of-the art generalization performance and data efficiency on Procgen [9]. However, we find that these factors do not always provide the same benefits for the task generalization benchmarks in Meta-World [74], indicating that transfer remains a challenge and may require different approaches than procedural generalization. Overall, we suggest that building generalizable agents requires moving beyond the single-task, model-free paradigm and towards self-supervised model-based agents that are trained in rich, procedural, multi-task environments.

A Survey of Zero-shot Generalisation in Deep Reinforcement Learning

“…World models which can adapt to changing dynamics will be more challenging, but Seo et al [148] give an initial example. Anand et al [156] is the first example investigating how well standard MBRL approaches generalise, and we look forward to seeing more work in this area.…”

“…A second under-explored area is model-based reinforcement learning (MBRL) for generalisation. Most methods surveyed here are model-free, with notable exceptions being [150,126,148,156]. Learning a world model and combining it with planning methods can enable stronger forms of generalisation, especially to novel reward functions (if the reward function is available during planning).…”

“…Another angle on better optimisation is the use of model-based RL (MBRL). Very little work has applied MBRL to generalisation benchmarks, with [156] being an initial example. Anand et al [156] apply MuZero Reanalyse [157,158], a SOTA MBRL method, to OpenAI Procgen [35], showing much-improved performance over SOTA model-free methods at much lower sample complexity.…”

“…Very little work has applied MBRL to generalisation benchmarks, with [156] being an initial example. Anand et al [156] apply MuZero Reanalyse [157,158], a SOTA MBRL method, to OpenAI Procgen [35], showing much-improved performance over SOTA model-free methods at much lower sample complexity. This shows the potential of using MBRL to improve generalisation to varying state and observations.…”

On the Effectiveness of Fine-tuning Versus Meta-reinforcement Learning

“…Prior work has explored improving generalization by training on a large amount of levels [13], adding dataaugmentation to visual inputs [21,22], knowledge transfer during training [23], and self-supervised world models [24]. Beyond zero-shot generalization, few prior works use Procgen to evaluate meta-RL, with one exception being Alver et al [25], which showed that RL 2 failed to generalize on simplified Procgen games.…”