Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

MAMBA: an Effective World Model Approach for Meta-Reinforcement Learning

Zohar Rimon, Tom Jurgenson, Orr Krupnik, Gilad Adler, Aviv Tamar

TL;DR

This work proposes a new model-based approach to meta-RL, based on elements from existing state-of-the-art model-based and meta-RL methods, and demonstrates the effectiveness of this approach on common meta-RL benchmark domains, attaining greater return with better sample efficiency while requiring very little hyperparameter tuning.

Abstract

Meta-reinforcement learning (meta-RL) is a promising framework for tackling challenging domains requiring efficient exploration. Existing meta-RL algorithms are characterized by low sample efficiency, and mostly focus on low-dimensional task distributions. In parallel, model-based RL methods have been successful in solving partially observable MDPs, of which meta-RL is a special case. In this work, we leverage this success and propose a new model-based approach to meta-RL, based on elements from existing state-of-the-art model-based and meta-RL methods. We demonstrate the effectiveness of our approach on common meta-RL benchmark domains, attaining greater return with better sample efficiency (up to $15\times$) while requiring very little hyperparameter tuning. In addition, we validate our approach on a slate of more challenging, higher-dimensional domains, taking a step towards real-world generalizing agents.

MAMBA: an Effective World Model Approach for Meta-Reinforcement Learning

TL;DR

This work proposes a new model-based approach to meta-RL, based on elements from existing state-of-the-art model-based and meta-RL methods, and demonstrates the effectiveness of this approach on common meta-RL benchmark domains, attaining greater return with better sample efficiency while requiring very little hyperparameter tuning.

Abstract

Meta-reinforcement learning (meta-RL) is a promising framework for tackling challenging domains requiring efficient exploration. Existing meta-RL algorithms are characterized by low sample efficiency, and mostly focus on low-dimensional task distributions. In parallel, model-based RL methods have been successful in solving partially observable MDPs, of which meta-RL is a special case. In this work, we leverage this success and propose a new model-based approach to meta-RL, based on elements from existing state-of-the-art model-based and meta-RL methods. We demonstrate the effectiveness of our approach on common meta-RL benchmark domains, attaining greater return with better sample efficiency (up to ) while requiring very little hyperparameter tuning. In addition, we validate our approach on a slate of more challenging, higher-dimensional domains, taking a step towards real-world generalizing agents.
Paper Structure (47 sections, 3 theorems, 10 equations, 4 figures, 7 tables)

This paper contains 47 sections, 3 theorems, 10 equations, 4 figures, 7 tables.

Table of Contents

  1. Introduction
  2. Background
  3. Problem Formulation
  4. Context-Based Meta-RL and Model-Based Algorithms for POMDP
  5. Method
  6. Comparing VariBAD and Dreamer
  7. Decomposable Task Distributions
  8. PAC Bounds
  9. Bounds on Bayes-suboptimality for DTD:
  10. Decomposability and Local Reconstruction
  11. MAMBA: MetA-RL Model-Based Algorithm
  12. Experiments
  13. Environments:
  14. MAMBA is an effective meta-RL algorithm:
  15. Robustness of MAMBA:
  16. ...and 32 more sections

Key Result

Theorem 1

(informal) Under mild assumptions, there exists a density estimator $\hat{f}$ which maps $N$ samples from a decomposable task distribution to a distribution over the task space, such that with probability at least $1-1/N$:

Figures (4)

  • Figure 1: Behaviors learned by MAMBA in various domains. In all of the figures, the first, second and third episodes are shown in red, green and blue, respectively. From left to right: Point Robot Navigation, Escape Room, Rooms-$N$ and Reacher-$N$ (plotted point represents the end-effector; see Sec. \ref{['sec:experiments']} for details). As the environments have sparse rewards, MAMBA evidently learns near-Bayes-optimal behavior: it explores the environment in the first episode and exploits the information in the subsequent ones. For additional visualizations see https://sites.google.com/view/mamba-iclr2024.
  • Figure 2: Reward prediction at current location (green), state identifier (blue), and goal state indication (red) through time on the Rooms-4 environment. Top: vanilla VariBAD (global reconstruction), bottom: VariBAD with local reconstruction. At the highlighted time steps (indicated by arrows) we plot the reconstructed reward for every grid cell, where positive reward prediction is represented as green, and the agent trajectory is drawn in blue.
  • Figure 3: Returns in Reacher-1: Left: sample efficiency -- meta-episode returns against the number of frames played. Since VariBAD and HyperX took 15$\times$ frames to converge we show their limit instead of a full plot. Right: time efficiency -- meta-episode returns against training time (hours).
  • Figure 4: Sub-episode Returns of MAMBA, VariBAD, HyperX and Dreamer-Vanilla in Reacher-1 (the meta-episode contains three sub episodes). Return of episode 0, 1, and 2, are left, middle, and right respectively. The gap between MAMBA and Dreamer demonstrates that MAMBA retains information from previous sub-episodes better than Dreamer.

Theorems & Definitions (4)

  • Theorem 1
  • Definition 1
  • Theorem 2
  • Theorem 3