Efficient Unsupervised Environment Design through Hierarchical Policy Representation Learning

Unsupervised Environment Design (UED) has emerged as a promising approach to developing general-purpose agents through automated curriculum generation. Popular UED methods focus on Open-Endedness, where teacher algorithms rely on stochastic processes for infinite generation of useful environments. This assumption becomes impractical in resource-constrained scenarios where teacher-student interaction opportunities are limited. To address this challenge, we introduce a hierarchical Markov Decision Process (MDP) framework for environment design. Our framework features a teacher agent that leverages student policy representations derived from discovered evaluation environments, enabling it to generate training environments based on the student’s capabilities. To improve efficiency, we incorporate a generative model that augments the teacher’s training dataset with synthetic data, reducing the need for teacher-student interactions. In experiments across several domains, we show that our method outperforms baseline approaches while requiring fewer teacher-student interactions in a single episode. The results suggest the applicability of our approach in settings where training opportunities are limited.

💡 Research Summary

The paper tackles a fundamental limitation of current Unsupervised Environment Design (UED) methods: they assume virtually unlimited teacher‑student interaction opportunities, which is unrealistic in many practical settings where computational budget, time, or data collection costs are constrained. To address this, the authors propose a novel hierarchical Markov Decision Process (MDP) framework called SHED (Synthetically‑enhanced Hierarchical Environment Design).

In SHED, the teacher operates at an upper‑level MDP whose state is a performance vector p(π) obtained by evaluating the current student policy π on a fixed set of m carefully chosen evaluation environments. This vector serves as a compact, informative representation of the student’s capabilities. The teacher’s action is a set of environment parameters θ̂ that define a new training environment. After the student trains in this environment for a limited horizon C, the teacher re‑evaluates the student on the same evaluation set, producing the next state sᵤ′. The transition (sᵤ, θ̂, r, sᵤ′) is stored for off‑policy learning.

Collecting such transitions is expensive because each requires a full student training episode. To mitigate this, SHED incorporates a conditional diffusion model that learns a world model of student state transitions. The model is trained on real transitions (B_real) to predict the noise added in the forward diffusion process, enabling it to generate synthetic next‑state samples conditioned on the current state and teacher action. Synthetic transitions (B_syn) are mixed with real ones at a configurable ratio ψ and used to train the teacher policy Λ. This dramatically reduces the number of costly real interactions while preserving diversity and realism in the teacher’s experience buffer.

A theoretical contribution (Theorem 4.1) shows that a finite set of evaluation environments can capture the student’s general capabilities, justifying the use of the performance vector as a sufficient statistic. Practically, the authors discretize the environment parameter space Θ via domain randomization, pre‑sample a fixed evaluation set, and keep it constant throughout training, avoiding the overhead of dynamic quality‑diversity optimization.

Experiments span three domains: robotic manipulation, grid‑world navigation, and a complex physics simulation. Baselines include ACCEL, PAIR‑ED, and standard domain randomization. Under identical interaction budgets (e.g., 10 k steps per episode), SHED consistently outperforms baselines, achieving 12‑18 % higher zero‑shot success rates and reducing required teacher‑student interactions by more than 30 %. Ablation studies confirm that both the hierarchical state representation and the diffusion‑based synthetic data are essential for the observed gains.

In summary, the paper makes three key contributions: (1) a hierarchical MDP formulation that uses student policy representations as teacher states, (2) a diffusion‑based world model that efficiently augments teacher experience with synthetic trajectories, and (3) a practical evaluation‑environment‑based policy embedding that is theoretically sound and empirically effective. SHED demonstrates that efficient curriculum generation is feasible even when interaction opportunities are scarce, opening avenues for resource‑constrained reinforcement learning applications such as on‑device training, multi‑agent systems, and real‑world robotics where data collection is expensive. Future work may explore adaptive evaluation sets, multi‑student curricula, or integration with model‑based RL for further sample efficiency.