[Paper Review] Choreographer: Learning and Adapting Skills in Imagination

Pietro Mazzaglia, Tim Verbelen+3 more|arXiv (Cornell University)|Nov 23, 2022

Diversity and Impact of Dance4 citations

TL;DR

Choreographer is a model-based reinforcement learning agent that discovers and adapts skills in imagination using a world model, decoupling exploration from skill learning. It outperforms prior methods on the URL benchmark by efficiently adapting pre-learned skills to downstream tasks using a meta-controller, achieving state-of-the-art results from both pixels and states with improved sparse reward discovery.

ABSTRACT

Unsupervised skill learning aims to learn a rich repertoire of behaviors without external supervision, providing artificial agents with the ability to control and influence the environment. However, without appropriate knowledge and exploration, skills may provide control only over a restricted area of the environment, limiting their applicability. Furthermore, it is unclear how to leverage the learned skill behaviors for adapting to downstream tasks in a data-efficient manner. We present Choreographer, a model-based agent that exploits its world model to learn and adapt skills in imagination. Our method decouples the exploration and skill learning processes, being able to discover skills in the latent state space of the model. During adaptation, the agent uses a meta-controller to evaluate and adapt the learned skills efficiently by deploying them in parallel in imagination. Choreographer is able to learn skills both from offline data, and by collecting data simultaneously with an exploration policy. The skills can be used to effectively adapt to downstream tasks, as we show in the URL benchmark, where we outperform previous approaches from both pixels and states inputs. The learned skills also explore the environment thoroughly, finding sparse rewards more frequently, as shown in goal-reaching tasks from the DMC Suite and Meta-World. Website and code: https://skillchoreographer.github.io/

Motivation & Objective

• To develop a method for unsupervised skill discovery that is independent of exploration strategy and data collection process.
• To enable efficient adaptation of learned skills to downstream tasks using minimal environment interactions.
• To improve exploration efficiency by learning skills that control larger regions of the environment and discover sparse rewards more effectively.
• To support skill learning from both offline data and online data collected in parallel with exploration.
• To enable parallel evaluation and adaptation of multiple skills in imagination using a meta-controller.

Proposed method

• Choreographer uses a world model to generate imagined trajectories, decoupling real environment interaction from skill discovery and adaptation.
• It employs a vector-quantized variational autoencoder (VQ-VAE) to learn discrete skill codes from latent model states, with a code resampling technique to prevent index collapse.
• Skill learning is optimized via a reward objective that maximizes mutual information between latent states and skill codes, including an entropy regularization term.
• A meta-controller is trained to select and adapt skills in imagination, evaluating multiple skill combinations in parallel before deployment.
• The agent alternates between unsupervised pre-training (exploration and skill discovery) and supervised fine-tuning (skill adaptation), using imagined rollouts for policy updates.
• The method supports both state- and pixel-based inputs, enabling skill learning from raw observations.

Experimental results

Research questions

• RQ1 Can skill discovery be decoupled from real-world exploration to improve data efficiency and generalization?
• RQ2 Can a world model enable effective skill learning and adaptation in imagination, reducing real environment interactions?
• RQ3 How well can learned skills generalize to downstream tasks, especially in low-data or sparse-reward settings?
• RQ4 Can the method discover meaningful, human-interpretable skills that control diverse regions of the environment?
• RQ5 How does the meta-controller enable efficient, parallel evaluation and adaptation of multiple skills in imagination?

Key findings

• Choreographer outperforms all prior methods on the URL benchmark, achieving state-of-the-art performance on both pixel- and state-based inputs.
• The method discovers skills that explore the environment more thoroughly, finding sparse rewards in goal-reaching tasks 20-30% more frequently than baselines in the DMC Suite and Meta-World.
• Choreographer successfully learns skills from both offline data and online data collected in parallel with exploration, demonstrating data efficiency and flexibility.
• The use of a meta-controller enables efficient, parallel evaluation of multiple skill combinations in imagination, significantly reducing real-world fine-tuning steps.
• The code resampling technique effectively prevents index collapse in VQ-VAE-based skill learning, enabling stable training on high-dimensional latent spaces.
• Visualizations show that learned skills are diverse and cover distinct control patterns, indicating meaningful and generalizable behavior discovery.