source/primary/appendix/16_additional_domains.tex

\section{Additional Task Domains}\label{app:extra_domains}

\subsection{Neural Cellular Automata}

\paragraph{Task Domain} Neural Cellular Automata (NCAs) are a form of cellular automata in which updates to cells---a function of their neighbours' states---are parameterised and learned by a neural network \cite{mordvintsev2020growing}. This task is about designing the components that the NCA uses to perceive its neighbours and learn its update rule.


\paragraph{Implementation} The code for Neural Cellular Automata is adapted from CAX \cite{faldor2025cax}, a library which provides a unified JAX framework for implementing cellular automata.


\paragraph{Editable modules} We include five editable modules: (i) the perceive module, (ii) the update module, and (iii) the loss, (iv) optimiser and (v) train modules for training the update module.


\paragraph{Datasets} We support five datasets: two growing tasks \cite{mordvintsev2020growing}, a self-classifying MNIST task \cite{randazzo2020self}, matrix operations \cite{bena2025path} and MNIST inpainting \cite{tesfaldet2022attention}.

% -=-=-=-=-=-=-=-=-=-=-=-=-=-==-=-=-=-=-=-

\subsection{Offline RL}

\paragraph{Task Domain} Offline Reinforcement Learning (offline RL) is a class of algorithms that trains a policy from previously collected environment transitions, without any other additional environment interactions~\citep{levine2020offline}. Unlike in online RL, the agent cannot explore and collect new trajectories from the \textit{online} environment and must extract the policy from pre-existing datasets. Such datasets may contain suboptimal trajectories or insufficient state coverage~\citep{kumar2022should}. In practice, the real environment---the \textit{online} domain inaccessible during training---can be used to tune the algorithm's hyperparameters~\cite{jackson2025a}.

\paragraph{Implementation} We base our implementation on Revisited Behaviour Regularized Actor-Critic (ReBRAC)~\citep{tarasov2023revisiting}, a high-performing and hyperparameter-robust offline RL algorithm. Our JAX implementation is based on those by \citet{jackson2025a} and \citet{park2025flowqlearning}.

\paragraph{Editable Modules} We include five editable modules for the offline RL task domain: (i) the actor loss, (ii) the critic loss, (iii) the network architectures, (iv) the optimizer, and (v) the training loop.

\paragraph{Datasets} We support all single-task, reward-labeled datasets from OGBench~\citep{ogbench_park2025}, a currently unsaturated offline RL benchmark. This encompasses several robot morphologies across various locomotion and manipulation tasks.

% -=-=-=-=-=-=-=-=-=-=-=-=-=-==-=-=-=-=-=-

\subsection{On-Policy Multi-Agent RL}

\paragraph{Task Domain} 
On-policy multi-agent reinforcement learning (MARL) refers to a class of multi-agent methods in which each agent updates its policy using an on-policy RL algorithm, thereby relying solely on trajectories generated by their current policy.

\paragraph{Implementation} 
The implementation of our on-policy MARL task is adapted from the Independent PPO (IPPO) algorithm \cite{ippo}, as provided in the JaxMARL library \cite{flair2024jaxmarl}.


\paragraph{Editable modules} 
We include six editable modules in On-Policy MARL: the advantage estimation and critic target computation, the loss function, the optimiser, the network architecture, the activation function used by the network, and the training loop.

\paragraph{Additional Backends}
In addition to the default, we support a recurrent architecture. In this case, the training loop must handle the recurrent state produced by the agent’s network.

\paragraph{Datasets}
We support 17 environments from the JaxMARL library \cite{flair2024jaxmarl}, comprising 5 multi-agent Brax tasks, 11 tasks from SMACv2 \cite{ellis2023smacv2}, and the MPE Spread environment \cite{mpe}.


% -=-=-=-=-=-=-=-=-=-=-=-=-=-==-=-=-=-=-=-


\subsection{Trajectory Prediction}

\paragraph{Task Domain} 
Trajectory prediction is the task of forecasting the future positions of traffic participants, such as vehicles, pedestrians, and cyclists, given their observed past motion and surrounding road context \cite{huang2022survey}.
Effective trajectory prediction must account for the multi-modal nature of human behaviour; at any moment, an agent may accelerate, brake, turn, or change lanes, which can lead to multiple plausible futures.
The objective is to produce a set of $K$ diverse future trajectories for a target agent, each with an associated probability, that collectively cover the space of likely outcomes.
Models are evaluated on how closely their best predictions match the ground truth using minimum Average Displacement Error (minADE), minimum Final Displacement Error (minFDE), miss rate, and Brier-minFDE. 
The ability to anticipate these possibilities is critical for safe motion planning in self-driving systems.

\paragraph{Implementation} 
The base code for Trajectory Prediction is adapted from UniTraj \citep{feng2024unitraj}, a unified framework for scalable vehicle trajectory prediction, and AutoBot \citep{girgis2022autobot}, a latent-variable sequential set transformer for joint multi-agent motion prediction.
The implementation provides a standardised pipeline for encoding past agent trajectories and road geometry via attention-based modules, decoding multi-modal future trajectories with uncertainty estimates, and evaluating predictions against ground truth.

\paragraph{Editable modules} 
We include four editable modules in Trajectory Prediction: (i) the loss function, which defines the multi-modal training objective balancing trajectory likelihood with mode diversity; (ii) the optimiser, which controls weight updates and learning rate scheduling; (iii) the network architecture, which encodes agent dynamics, social interactions, and map context before decoding future trajectories; and (iv) the training loop, which orchestrates data loading, model training, validation, and model selection.

\paragraph{Datasets} 
We support three large-scale autonomous driving datasets: Argoverse~2 \citep{wilson2023argoverse}, collected across six US cities (Austin, Detroit, Miami, Pittsburgh, Palo Alto, and Washington D.C.); nuScenes \citep{caesar2020nuscenes}, recorded in Boston and Singapore and covers diverse urban conditions; and the Waymo Open Motion Dataset \citep{ettinger2021large}, spanning six US cities (San Francisco, Phoenix, Mountain View, Los Angeles, Detroit, and Seattle).
Each dataset is preprocessed into a standardised format with 21 observed timesteps (2.1\,s at 10\,Hz), up to 32 surrounding agents, 128 map polylines, and a 60-timestep (6\,s) prediction horizon.
All three datasets contain 850 preprocessed scenarios, hosted on HuggingFace.
\newpage

\section{Total Task Counts} \label{app:counts}

Including the additional domains from Appendix \ref{app:extra_domains}, and a small number of extra modules in On-Policy RL (the target and network activation function), dramatically increases the number of tasks in DiscoGen's support. As such, in Table \ref{tab:extra_task_domains}, we provide a more complete set of task counts. This also demonstrates the power of the combinatoric explosion of expanding domains; simply adding 2 modules to On-Policy RL effectively quadrupled the number of tasks in its support, to over a billion.

\begin{table}[h!]
\centering
\caption{Overview of all domains and their number of supported tasks.\vspace{-5pt}}
\label{tab:extra_task_domains}
\setlength{\tabcolsep}{8pt}
\begin{tabularx}{0.6\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  & 1 & 3  & 1 &  85,176 \\
Neural Cellular Automata       & 5 & 5  & 1 & 33,480 \\
Off-Policy RL                   & 7 & 4  & 1 & 38,100 \\
Offline RL                   & 5 & 10  & 1 & 10,602,372 \\
On-Policy MARL               & 6 & 17  & 2 & 97,431,783,120 \\
On-Policy RL               & 6 & 13 & 3 & 1,789,383,960 \\
Trajectory Prediction          & 4 & 3 & 3 & 1,080 \\
Unsupervised Environment Design & 3 & 4  & 1 & 2,100 \\
\bottomrule
Total & & & & 99,299,115,384 \\
Median & & & & 59,622
\end{tabularx}
\vspace{-4pt}

\end{table}