Neural Operators for Multi-Task Control and Adaptation

Abstract

Neural operator methods have emerged as powerful tools for learning mappings between infinite-dimensional function spaces, yet their potential in optimal control remains largely unexplored. We focus on multi-task control problems, whose solution is a mapping from task description (e.g., cost or dynamics functions) to optimal control law (e.g., feedback policy). We approximate these solution operators using a permutation-invariant neural operator architecture. Across a range of parametric optimal control environments and a locomotion benchmark, a single operator trained via behavioral cloning accurately approximates the solution operator and generalizes to unseen tasks, out-of-distribution settings, and varying amounts of task observations. We further show that the branch-trunk structure of our neural operator architecture enables efficient and flexible adaptation to new tasks. We develop structured adaptation strategies ranging from lightweight updates to full-network fine-tuning, achieving strong performance across different data and compute settings. Finally, we introduce meta-trained operator variants that optimize the initialization for few-shot adaptation. These methods enable rapid task adaptation with limited data and consistently outperform a popular meta-learning baseline. Together, our results demonstrate that neural operators provide a unified and efficient framework for multi-task control and adaptation.

Figures (10)

• Figure 1: A point-to-point multi-task control problem. Left: A single point-to-point task. Middle: A series of tasks represented by their cost functions $c_i$. Right: The corresponding optimal policies $\pi^*_i$ where the small arrows correspond to the control outputs at different x,y positions. This is naturally modeled as a mapping between function spaces.
• Figure 2: DeepONet/SetONet architecture: Here we show the mapping $\mathcal{T}_\theta[\ell_i] \rightarrow \hat{\pi}_i$, with pointwise evaluations of $\ell_i$ (in red) and of $\hat{\pi}_i$ at the point $\mathbf{y}$. The branch network maps sensor locations $(x, u)$ of a cost function $\ell(x, u; {\boldsymbol{\phi}})$ to task-dependent coefficients $\{c_k(\ell)\}_{k=1}^p$. The red points indicate the pointwise samples of $\ell_i$. The trunk maps query locations $\mathbf{y} = (x, t)$ to a set of learned basis functions $\{b_k(\mathbf{y})\}_{k=1}^p$. Their inner product yields the predicted control output at the query location $\mathbf{y}$
• Figure 3: Overview of the meta-training procedure. The inner loop adapts the parameters $\theta^k$ to each task $i$ using the support set $\mathcal{D}_i^{\mathrm{tr}}$. SetONet-Meta updates only the branch coefficients $\theta_{\mathrm{branch}}$, while SetONet-Meta-Full updates all parameters $\theta$. The adapted parameters $\theta_i'$ from a batch of $B$ tasks are evaluated on held-out query sets $\mathcal{D}_i^{\mathrm{eval}}$, and the outer loop computes the meta-gradient to update the shared initialization $\theta^{k+1}$.
• Figure 4: Operator fitting results across three environments. Each row group shows two control dimensions for a given environment, with columns displaying predictions on two representative tasks. We then show the corresponding state-space rollouts for the two tasks (T1, T2) along with 3 more (T3-T5). Solid green lines denote expert demonstrations, solid blue are SetONet predictions at the expert state locations, and dashed lines denote model-rollouts using the learned operator policy.
• Figure 5: Task resolution invariance across all four control environments. Lines show median relative $L^2$ error over held-out tasks; shaded regions indicate the inter-quartile range. Blue tick marks denote context sizes seen during training; red tick marks denote sizes not seen during training. The $x$-axis for the first three environments is the number of context samples; for Obstacle Avoidance it is the number of obstacles. The model was trained on obstacle configurations with 2, 4, and 6 obstacles.