Choreographer: Learning and Adapting Skills in Imagination

TL;DR

Choreographer tackles data inefficiency in unsupervised RL by learning, discovering, and adapting diverse skills inside imagination through a world model. It decouples exploration from skill learning using a VQ-VAE codebook to map model states to discrete skill codes and optimizes a mutual information objective $I(\tau_s; z)$, with rewards combining entropy and code-distance terms $r_{\text{skill}} = r_{\text{ent}} + r_{\text{code}}$, all trained via imagination-based rollouts. Adaptation to downstream tasks is achieved with a meta-controller that selects skill codes in imagination, enabling efficient fine-tuning guided by a reward predictor. Empirical results on the URL benchmark and sparse-reward scenarios (e.g., Jaco, MetaWorld) show state-of-the-art data efficiency and improved exploration, with code and resources released for reproducibility.

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/

Figures (20)

• Figure 1: Unsupervised Reinforcement Learning. The agent should effectively leverage the unsupervised phase, consisting of the data collection and the pre-training (PT) stages, to efficiently adapt during the supervised phase, where the agent is fine-tuned (FT) for a downstream task.
• Figure 1: Sparse rewards. Finding goals in the environment with sparse rewards using the pre-trained behaviors. The table shows the fraction of goals found by each method per task. (3 seeds)
• Figure 2: Choreographer. (a) The agent leverages representation learning to learn a codebook of skill vectors, summarizing the model state space. (b) Skill policies are learned in imagination, maximizing the mutual information between latent trajectories of model states and skill codes. (c) The meta-controller selects the skills to apply for downstream tasks, evaluating and adapting skill policies in imagination by hallucinating the trajectories of states and rewards they would obtain.
• Figure 3: Offline state-based URLB. On the left, Choreographer (ours), pre-trained with offline exploratory data, performs best against baselines, both in terms of IQM and Optimality Gap. On the right, mean and standard deviations across the different domains of URLB are detailed (10 seeds).
• Figure 4: Pixel-based URLB. Performance of Choreographer (ours) as a function of pre-training steps. Scores are asymptotically normalized and averaged across tasks for each method. (10 seeds)