source/primary/sections/5_discogen.tex

\vspace{-3pt}\section{DiscoGen} \label{sec:discogen}

% \vspace{-2pt}
% \subsection{The DiscoGen Task Framework} \label{sec:framework}
% \vspace{-2pt}

DiscoGen generates \textbf{tasks}; modular algorithm discovery problems consisting of an objective and meta-train and meta-test datasets. Each task is defined by seven components.


\paragraph{1. Task Domain} The task domain establishes which area of machine learning a task pertains to (e.g., \textit{On-Policy RL} or \textit{Image Classification}). It defines the initial codebase, agent objective, and which datasets and modules are available.

\paragraph{2. Editable Modules} For each task domain, we identify important building blocks, or \textit{modules}, of an algorithm that can be set as \textit{editable} or \textit{fixed}. Whereas \textit{fixed} modules use standard implementations (e.g., a conventional loss function or optimiser), \textit{editable} modules specify the interface for an agent's code only (i.e., inputs, outputs and generic function names). This reduces the bias of agents, while ensuring their implementation fits the rest of the codebase. There can be many modules per domain, combinatorially expanding the number of possible tasks that can be generated.

\paragraph{3. Meta-Training Datasets} Meta-training datasets are the problems which an agent can run experiments on during algorithm discovery. They are known to the agent.

\paragraph{4. Meta-Test Datasets} Meta-test datasets are used to test the agent's discovered algorithm \textit{after} the meta-loop has completed. They test how well algorithms transfer to held-out problems, and are not known to the agent.


\paragraph{5. (Optional) Backend} For some task domains, we implement additional \textit{backends} that provide new tasks without warranting their own domain. For example, in On-Policy RL, a user can specify whether to use a feed-forward or recurrent policy, providing a new search space for algorithms.

\paragraph{6. Evaluation Type} Each task can use one of a number of meta-learning objectives, which the ADA must optimise. Currently, DiscoGen supports three evaluation types for algorithms: its final score (maximise performance); the energy used to reach some proportion of the baseline's score (maximise efficiency); or the time taken to reach a proportion of the baseline's score (maximise speed), like \citet{zhao2025the} or \citet{kellerjordan2024moddednanogpt}. In this paper we focus on evaluating for performance only, due to resource constraints. Expanding DiscoGen to other evaluation types is trivial.

\paragraph{7. Initialisation} DiscoGen supports task generation using either `empty' initialisations, in which only the interface of each editable module is provided, or `baseline' initialisations, where editable modules start from full implementations of baseline algorithms (as in \citet{nathani2025mlgymnewframeworkbenchmark}). Our experiments use `empty' initialisations in this paper.

\subsection{An Example Task}
\begin{lstlisting}[language=Python, caption={An example task configuration. A task is defined by its task domain, meta-train/meta-test datasets, backend and modules.}, label={fig:example_config}]
task_domain = "OnPolicyRL"
meta_train = ["Breakout", "Freeway"]
meta_test = ["Asterix", "SpaceInvaders"]
backend = "recurrent"
change_optim = False
change_loss = True
change_networks = True
change_train = False
eval_type = "performance"
initialisation = "empty"
\end{lstlisting}

\vspace{-10pt}
We use Example \ref{fig:example_config} to demonstrate the task interface. For this on-policy RL task, the ADA will work in JAX \citep{bradbury_jax_2018} to maximise the performance of an RL agent. The majority of the codebase is \textit{fixed}: domain-specific code like wrapper functions and environment creation, as well as the optimiser (which uses Adam \citep{kingma_adam_2017}) and training loop modules. While an agent \textit{can} interact with these files, any changes will be overwritten prior to meta-testing to prevent evaluation hacking. 

The ADA must work with two \textit{editable} files: the loss, which starts as an empty function mapping inputs like collected data batches to a scalar loss; and the network architecture, which maps environment observations to an RL policy, value function and recurrent state and also includes no logic. The ADA can make arbitrary code edits, such as writing additional functions, importing missing packages, and filling the templates. The agent can \textit{run} inner-loop training and evaluation for both meta-train environments. To ensure security during agent execution, the ADA operates in a containerised environment. If one editable module is \textit{only} called by another, the ADA could also edit its interface, expanding the search-space of algorithms in DiscoGen tasks further.\vspace{-6pt}

\subsection{Procedural Task Generation} \label{sec:domains}
\vspace{-5pt}
DiscoGen is a procedural generator of ML tasks. As a procedural generator, it is designed for \textit{improving} ADAs, manually or automatically in an ADA optimisation loop (\cref{fig:discogen}). A task is specified by a small configuration file, like Example \ref{fig:example_config}, which takes the same role as the parameters in other PCG environments \citep{cobbe_leveraging_2020, team_open-ended_2021}. This can be randomly generated, user-specified, or sampled using autocurricula. The technical details of sampling tasks in DiscoGen are described in Appendix \ref{app:technical_implementation}.\vspace{-1pt}


DiscoGen creates tasks in a two-stage process to reduce the risk of meta-test leakage. The meta-train portion of the task is generated first; only \textit{after} the meta-loop is complete does DiscoGen create the meta-test codebase. As such, the agent is never given any details of the meta-test datasets.\vspace{-1pt}

PCG involves creating levels, or tasks, programmatically. DiscoGen is no different; the number of tasks for a domain is combinatorial with respect to how many modules and datasets it supports. Specifically, for a task domain with $m$ modules, $d$ datasets, and $b$ backends, and our currently supported $3$ evaluation types and $2$ task initialisations, \vspace{-5pt}
\begin{equation}
    N_{tasks} = 2 \cdot 3 \cdot b \cdot (2^m-1)\cdot\left(3^d-2^{(d+1)}+1\right).
\end{equation}\vspace{-19pt}

Derived in Appendix \ref{app:derivation}, this assumes at least one editable module and assigns each dataset to either meta-train, meta-test or unused (with at least one meta-train and meta-test).

\vspace{-5pt}
\subsection{Available Task Domains}\vspace{-3pt}

Table \ref{tab:task_domains} provides a snapshot of the domains in DiscoGen, which represent diverse fields ranging from simple applied problems to complex foundational ones, spanning supervised to reinforcement learning. We expect this collection to grow with open-source community contributions. We describe each domain, its modules and datasets in Appendix \ref{app:domains}. We describe a small number of additional domains and editable modules which are not analysed in this paper, due to computational limitations, in Appendix \ref{app:extra_domains}, and the associated full task count ($\sim 100B$) in Appendix \ref{app:counts}.
Creating new domains is simple, needing only mild adaptation of existing codebases. Due to imbalanced task counts, we recommend uniformly sampling domains to reduce bias; such functionality is supported within the DiscoGen library.\vspace{-6pt}
\begin{table}[h!]
\centering
\caption{Overview of domains and their number of supported tasks.\vspace{-5pt}}
\label{tab:task_domains}
\begin{small}
\setlength{\tabcolsep}{4pt}
\begin{tabularx}{1\columnwidth}{@{}Xcccr@{}}
\toprule
\textbf{Task Domain} & \textbf{$m$} & \textbf{$d$} & \textbf{$b$} & \textbf{$N_{tasks}$} \\
\midrule
Bayesian Optimisation          & 6 & 11 & 1 & 65,413,656 \\
Brain Speech Detection          & 3 & 7  & 1 & 81,144 \\
Computer Vision Classification  & 4 & 9  & 1 & 1,679,400 \\
Continual Learning             & 5 & 3  & 3 & 6,696 \\
Greenhouse Gas Prediction       & 2 & 4  & 1 & 900 \\
Language Modelling             & 3 & 4  & 2 & 4,200 \\
Model Unlearning\footnotemark  & 1 & 3  & 1 &  85,176 \\
Off-Policy RL                   & 7 & 4  & 1 & 38,100 \\
On-Policy RL                    & 4 & 13 & 3 & 426,043,800 \\
Unsupervised Environment Design & 3 & 4  & 1 & 2,100 \\
\bottomrule
Total & & & & 493,355,172 \\
Median & & & & 59,622
\end{tabularx}
\end{small}\vspace{-4pt}

\end{table}
\footnotetext{Model unlearning entails finetuning pretrained models. For $n$ models, $N_{tasks} = 2\cdot3\cdot b \cdot (2^m-1)\cdot\left((2n+1)^d-2(n+1)^{d}+1\right)$} 
\vspace{-12pt}


DiscoGen exhibits useful diversity across these millions of tasks; in \cref{sec:prompts}, we show that ADA optimisation improves as more DiscoGen tasks are experienced. The domains supported in DiscoGen span a range of machine learning fields, incorporate datasets of varying complexity and difficulty, and are built upon codebases of differing scales. Most task domains include unique modules, and when there is overlap, implementations are domain-specific. In Appendix \ref{app:onpolicyresults}, we explore how ADA performance changes in on-policy RL for all 15 possible module combinations. Our analysis reveals that increasing the number of editable modules lowers the ADA's success rate, but the increased flexibility raises its achievable performance ceiling. 

We further validate this diversity through \textit{rank correlation analysis} over the performance of different ADAs in DiscoBench (\cref{sec:discobench_results}), a subset of DiscoGen tasks, in Figures \ref{fig:all_data_corr} \& \ref{fig:split_corrs} (Appendix \ref{app:correlation}). Hierarchical clustering of the correlation matrix reveals distinct patterns; while correlation is often high between similar modules in different domains or different modules in the same domain, there are also anti-correlations where strong performance in one task implies poor performance in another, including within the same domain. The Fisher-Z transformed mean Spearman correlation is $\sim0.4$; high enough to indicate non-random signal, while sufficiently small to show low redundancy in DiscoGen. Notably, we also find that clustering patterns are \textit{distinct} between meta-train and meta-test, demonstrating how the \textit{same algorithm} can rank differently across datasets. 

Per-dataset analysis in each task reinforces the meaningfulness of DiscoGen's large task space. Considering Appendix \ref{app:discobench_results}, the ranking of the discovered algorithms changes across datasets within the \textit{same task}. This is intuitive, given prior literature suggests optimal algorithms differ between datasets or RL environments (e.g., in reinforcement learning \citep{goldie_can_2024, jackson2025a}, computer vision \citep{rodrigo_comprehensive_2024, takahashi_comparison_2024} or language modelling \citep{transformersssms2024dao, repeatafter2024jelassi}).
\vspace{-8pt}
\subsection{Advantages of DiscoGen}\label{sec:advantages}
\vspace{-3pt}
Individual task design in DiscoGen also overcomes the many flaws of previous task definitions raised in \cref{sec:flaws}.\vspace{-9pt}

\paragraph{Principled Evaluation \& Contamination Resistance} DiscoGen tasks clearly distinguish between meta-train and meta-test. Since DiscoGen is procedural, and there is no knowledge of the meta-test datasets in meta-training, the potential for test leakage is limited. Even as DiscoGen enters pre-training datasets, this is a step towards fairer evaluation, ensuring DiscoGen remains pertinent for a long time. Furthermore, DiscoGen supports different evaluation \textit{types}, enabling discovery for factors other than performance.\vspace{-8pt}

\paragraph{High Diversity} DiscoGen generates highly diverse tasks. As demonstrated in \cref{sec:domains}, DiscoGen supports a range of domains with different data structures, filesystem complexities and modules. Since DiscoGen tasks are parameterised combinatorially, its tasks represent a range of difficulties from ``\textit{implement a single module for one easy dataset}'' to ``\textit{implement many modules for many hard datasets}''.\vspace{-7pt}

\paragraph{Different Initialisations} Agents need not implement full codebases, and the initialisation of editable modules can be set to just the inputs and outputs of the module only (empty) or fully functioning baseline implementations. In essence, tasks can involve either `\textit{improvement from a baseline}' or `\textit{de novo discovery}'. This allows us to better analyse the biases elicited by ADAs, can be used to make tasks easier or harder by changing the level of guidance, and may enable improved creativity in research agents.\vspace{-8pt}

\paragraph{Ease of Adding Tasks} Beyond our currently implemented domains, adding many tasks to DiscoGen is significantly easier than for other suites. For similar effort to adding one task to, say, MLE-Bench \citep{chan_mle-bench_2025} or MLGymBench \citep{nathani2025mlgymnewframeworkbenchmark}, DiscoGen can gain potentially millions of new tasks in its support. When the base code for a task domain is complete, adding more tasks is even easier; isolating a new module effectively doubles the number of possible tasks, and adding a new dataset near-triples it.\vspace{-8pt}

\paragraph{Unsaturated Problems} Since tasks can be made more or less difficult by changing the module and dataset configurations (Appendix \ref{app:onpolicyresults}), DiscoGen tasks span a wide range of difficulties. We find that the more modules there are to implement, the harder the problem but the higher the potential ceiling. Additionally, almost all datasets currently in DiscoGen have yet to be solved by humans, let alone agents, and adding more, harder datasets is straightforward.\vspace{-7pt}