Beyond Fixed Tasks: Unsupervised Environment Design for Task-Level Pairs

AI summary

Overview

Research area: Reinforcement learning — specifically unsupervised environment design (UED), task-conditioned policy learning, formal task specification via reward machines, and automatic curriculum generation.

Technical level: Advanced. The paper assumes familiarity with POMDPs, UPOMDPs, regret-based curriculum learning, finite-state machines, graph neural networks, and actor-critic methods such as PPO.

Scope: The paper introduces ATLAS, a method that co-evolves tasks and levels into a single joint autocurriculum, and an accompanying evaluation suite that pairs reward-machine tasks with Minigrid levels where solvable task-level pairs are rare.

What This Paper Is About

Training a general agent requires both a task to follow and an environment (a level) in which to follow it. When tasks and levels are sampled independently, most resulting pairs are unsolvable — the paper reports that under its default independent sampling setup, only 2.7% of sampled problems per batch are solvable. Existing unsupervised environment design (UED) methods can generate curricula of increasingly hard levels, but only for a fixed task.

ATLAS extends UED so that tasks and levels are co-designed, producing a joint curriculum of problems that are both solvable and appropriately challenging for the current policy. The paper also contributes an evaluation suite (reward-machine tasks embedded in Minigrid levels, plus 150 hand-designed problems) built specifically to stress-test this low-solvability regime.

Key Contributions

1. Joint task-level autocurricula. ATLAS (Aligning Tasks and Levels for Autocurricula of Specifications) extends regret-based UED from level-only generation to simultaneous co-design of tasks (reward machines) and levels, ensuring training problems are solvable yet near the frontier of the agent's ability.

2. Structure-aware task mutations. The paper adds task mutations that exploit the reward machine's structure, plus hindsight edits derived from partial progress (rollouts ending in an intermediate RM state), enabling exploration of a broader problem neighbourhood than curation alone. This is instantiated with ACCEL and PLR⊥.

3. An evaluation suite for rarely solvable settings. Reward machines encoding BabyAI-style instructions are embedded in Minigrid levels, with samplers supporting independent, level-conditioned, and task-conditioned problem generation. The reward machine alphabet has 889 propositions, which the authors state is almost two orders of magnitude larger than alphabets in prior work. The suite includes 150 hand-designed problems targeting long-term planning, implicit subgoals, and exploration.

4. A problem-conditioned policy architecture. An actor-critic policy conditioned on both environment observations and reward machine state, where the RM is encoded by a graph neural network (rather than a fixed embedding such as a one-hot state index) to handle varying RM topologies and edge labels. Training uses PPO.

Main Findings

• Regret-based methods dominate domain randomization. Under independent problem sampling, DR "solves almost no problems," while PLR⊥ and ACCEL (ATLAS instantiations) substantially outperform it on CVaR and on zero-shot performance over the hand-designed test set. The authors attribute this to the fact that independent sampling yields only 2.7% solvable problems per batch, making task-level co-design essential.

• Mutations help, especially in the worst case. ACCEL variants consistently outperform PLR⊥ across most α values in the CVaR metric, particularly for low α (worst-case problems). At α = 100% (average performance), ACCEL and PLR⊥ perform similarly and both substantially outperform DR.

• Mutations accelerate convergence. On the hand-designed evaluation set, ACCEL variants and PLR⊥ reach comparable final results, but ACCEL converges faster in early training.

• Mutations alone can build the curriculum. ACCEL-0, which starts from single-room levels with one object and reward machines with one transition, achieves strong performance — showing problem complexity can emerge purely through targeted mutations without initial problem diversity.

• Curricula emerge along all dimensions. Buffer problems for PLR⊥ and ACCEL start simple (few rooms, objects, RM states) and grow in complexity over training. ACCEL variants reach higher RM state counts than PLR⊥, and both prioritise dense six-room layouts near the end of training.

• Solvable buffer fraction grows to near 100%. The fraction of solvable problems in the buffer steadily grows, eventually nearing 100%, with faster growth in ACCEL than PLR⊥.

• Level-conditioned sampling lets DR catch up. An ablation with level-conditioned sampling raises solvability to 83.4%. As hypothesised, DR improves substantially, matching the CVaR and test performance of PLR⊥ and ACCEL variants, reported as approximately 91.6%.

• Robustness when solvable problems are scarce. In that same level-conditioned ablation, PLR⊥ and ACCEL see only marginal gains (approximately 1–3%) over the independent setting.

• Task complexity has a cost. Replacing the sequential task sampler with a random walk-based sampler restricted to directed acyclic graphs, solve rates on the hand-designed evaluation set drop by approximately 35% compared to the sequential setting, though PLR⊥ still substantially outperforms DR.

• Joint mutations are necessary. Disabling either level edits or task edits significantly reduces performance; combining both yields the best results.

• Sequence length sensitivity. Short mutation sequences severely hinder performance: single-edit sequences cause an approximately 40% drop and three-edit sequences an approximately 15% drop, while long sequences (20 edits) have minimal effect (approximately 2% variation).

• Hindsight edits are rare. Task edits and level edits contribute roughly equally to the curriculum, but hindsight edits are uncommon, which the authors attribute to limited applicability — they require rollouts to end in an intermediate RM state and can only appear as the first edit.

• Average edits per problem increases over training, indicating the buffer composition shifts from primarily randomly generated problems to predominantly mutated ones.

Methodology in Plain English

The authors frame the problem as a game between a teacher that proposes problems and a student policy that trains on them. Each problem is a pair: a task (an instruction) and a level (an environment instance). Tasks are written as reward machines — small finite-state machines whose edges are labelled by propositions like "in front of a ball" and whose transitions carry rewards. Levels are Minigrid environments containing objects (keys, squares, balls, doors) in various colours, arranged across rooms connected by doors.

Rather than sampling tasks and levels uniformly, ATLAS keeps a buffer of high-regret problems — problems where the gap between the best achievable return and the current policy's return is large. This directs training toward solvable problems at the edge of what the agent can currently do, while screening out unsolvable ones. Regret is approximated following MaxMC-style scoring, and the paper builds on Prioritized Level Replay (PLR), using its robust variant PLR⊥ and ACCEL.

ACCEL adds mutation: before a problem enters the buffer, it is edited a number of times, where the number of edits is sampled from a range and each edit type is chosen uniformly at random. Three edit families exist: level edits (move the agent, add or remove rooms, add, remove, replace or move objects), task edits (switch a proposition, add or remove a state), and hindsight edits (use partial progress in a rollout to create a simpler "preceding" subproblem or a "succeeding" continuation problem).

The policy is an actor-critic network trained with PPO. Observations go through a convolutional neural network; the reward machine goes through a graph neural network that produces embeddings for all RM states; the current state's embedding is concatenated with the CNN features and passed through a recurrent network to capture history.

Evaluation uses two metrics, averaged over five seeds with 95% confidence intervals, with each problem evaluated 10 times per seed: CVaR (solve rate on the worst-performing α% of a large sampled problem set, measuring robustness) and the inter-quartile mean (IQM) of solve rates on the hand-designed evaluation set.

Default training levels draw a number of rooms from {1, 2, 4, 6}, a grid-dependent number of objects, the objects themselves, and the agent position. Tasks use the sequential sampler with path lengths randomly chosen between 1 and 5, and sparse rewards (1 on transitions to the accepting state, 0 otherwise).

Why This Matters

Impact on research. The paper identifies and targets an understudied regime — task-level pairs that are rarely solvable — where domain randomization collapses. It extends UED, previously limited to a fixed task, to joint task-and-level curriculum generation, and provides a benchmark suite (889 propositions, 150 hand-designed problems) for measuring progress. It also offers a hypothesis for why prior formal-language conditioning work with DR succeeds: their domains may ensure high solvability because all propositions are observable across levels.

Real-world applications:

• Training robots or embodied agents that must follow varied instructions in varied environments, where most instruction-environment combinations are infeasible.
• Automated generation of training curricula for industrial control or simulation environments without hand-authoring every scenario.
• Generating evaluation suites that stress-test agents on worst-case problems rather than average ones.
• Curriculum design for hierarchical or multi-step procedures (for example, multi-stage recipes or inspection routines) expressed as finite-state specifications.

Industry relevance. Any organisation training reinforcement learning agents where task specification and environment configuration are decoupled benefits from a method that filters out unsolvable combinations automatically. The finding that DR is adequate when solvability is high but fails when it is low gives practitioners a clear diagnostic: measure solvability in your sampling distribution first. The JAX-based implementation supporting CPUs, GPUs, and TPUs also lowers the barrier to adoption at scale.

Future Directions

• Extend to other temporal formalisms. The authors propose applying joint curriculum generation to linear temporal logic, programs, or hierarchies of reward machines (HRMs), noting the codebase already supports HRMs but the paper focuses on flat reward machines.
• Better scoring functions and mutation operators. Developing scoring functions and mutations that more fully exploit task structure could yield more effective curricula.
• More general hindsight edits. The paper notes that hindsight edits are rarely applied because they require rollouts to end in intermediate RM states and can only be the first edit in a sequence; exploring more general forms is flagged as future work.
• Scaling to richer domains. The authors propose testing joint curricula in domains such as Craftax.
• Addressing known UED limitations. The paper suggests irreducible regret floors may emerge under any scoring function, and that task-level variations offer a way to evaluate limitations previously identified for level-only UED (lack of diversity among high-regret levels, irreducible regret, and regret approximations that correlate with success rate rather than true regret).

Target Audience

Researchers and advanced practitioners in reinforcement learning, particularly those working on curriculum learning, unsupervised environment design, generalist agents, task-conditioned policies, and formal specification languages (reward machines, temporal logic, automata). Also relevant to engineers building training infrastructure for multi-task embodied agents, and to researchers studying generalisation and robustness benchmarks where solvable task-environment pairs are rare. The paper is not beginner-friendly: it assumes working knowledge of POMDPs, regret-based UED, Markov decision processes, graph neural networks, and PPO.

Authors’ abstract

Training general agents to follow complex instructions (tasks) in intricate environments (levels) remains a core challenge in reinforcement learning. Random sampling of task-level pairs often produces unsolvable combinations, highlighting the need to co-design tasks and levels. While unsupervised environment design (UED) has proven effective at automatically designing level curricula, prior work has only considered a fixed task. We present ATLAS (Aligning Tasks and Levels for Autocurricula of Specifications), a novel method that generates joint autocurricula over tasks and levels. Our approach builds upon UED to automatically produce solvable yet challenging task-level pairs for policy training. To evaluate ATLAS and drive progress in the field, we introduce an evaluation suite that models tasks as reward machines in Minigrid levels. Experiments demonstrate that ATLAS vastly outperforms random sampling approaches, particularly when sampling solvable pairs is unlikely. We further show that mutations leveraging the structure of both tasks and levels accelerate convergence to performant policies.