cs.LGApr 3, 2026
SaveNeural Operators for Multi-Task Control and Adaptation
Organizations: The University of Texas at Austin
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 & tables
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 . Right: The corresponding optimal policies 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 , with pointwise evaluations of (in red) and of at the point . The branch network maps sensor locations of a cost function to task-dependent coefficients . The red points indicate the pointwise samples of . The trunk maps query locations to a set of learned basis functions . Their inner product yields the predicted control output at the query location
Figure 3: Overview of the meta-training procedure. The inner loop adapts the parameters to each task using the support set . SetONet-Meta updates only the branch coefficients , while SetONet-Meta-Full updates all parameters . The adapted parameters from a batch of tasks are evaluated on held-out query sets , and the outer loop computes the meta-gradient to update the shared initialization .
| Environment | Expert | Tasks | Traj. | Task parameters | |
|---|---|---|---|---|---|
| P2P-Cost | 4/2 | LQR | 500 | 100 | |
| P2P-Cost-Small | 4/2 | LQR | 50 | 10 | |
| P2P-Dynamics | 4/2 | iLQR | 100 | 100 | |
| Quadrotor | 6/2 | iLQR | 100 | 20 | |
| Obstacle | 4/2 | NLP | 500 | 60 | |
| HalfCheetah-v3 | 17/6 | SAC | 53 | 100 | mass, limb, joints, friction |
Table 1: Summary of experimental environments. All OCP environments use the same SetONet architecture and training procedure; only the context and query definitions change across environments. The iMuJoCo environment encodes tasks through state–action transition data rather than analytical function evaluations. “Tasks” denotes the number of distinct task configurations and “Traj.” the number of expert trajectories per task.
| Method | P2P-Cost | P2P-Small | P2P-Dyn. | Quadrotor | Obstacle |
|---|---|---|---|---|---|
| 0 steps (zero-shot) | |||||
| Pretrained | .048 | .101 | .179 | .063 | .238 |
| MAML | .583 | .920 | .802 | .549 | 1.00 |
| SetONet-Meta | .075 | .084 | .276 | .185 | .287 |
| SetONet-Meta-Full | .118 | .130 | .195 | .148 | .395 |
| 1 gradient step | |||||
Table 2: Evaluation of the SetONet model across four OCP environments and a reduced-data variant (P2P-Small), using a single expert trajectory for demonstration and 32 trajectories for evaluation. The top section reports zero-shot predictions from the pretrained and meta-trained operators with no gradient updates. The lower sections report performance after 1 and 25 gradient steps of adaptation. Bold indicates the best result per method and column. “–” indicates divergence. The pretrained operator achieves low zero-shot error on P2P-Cost and Quadrotor, while dynamics-varying and obstacle environments benefit from adaptation, particularly via meta-trained initializations. All errors are reported as relative 19
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 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 -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.
Figure 6: Per-task comparison of MAML against four SetONet-based methods across three OCP environments after 25 gradient steps using 10 expert trajectories for adapting and evaluation. Each point represents a single held-out task, with the -axis showing the method’s relative error and the -axis showing MAML’s error on the same task. Points above the diagonal ( ) indicate tasks where MAML performs worse. MAML is consistently dominated across nearly all tasks in P2P-Cost and P2P-Dynamics, while in the Quadrotor environment the gap narrows, with some tasks falling near the diagonal.
Figure 7: Cost-based fine-tuning across two environments. (a) Out-of-distribution P2P-Cost tasks with both goal locations and start position outside the training region (gray). The distance is in meters. The pretrained SetONet fails to reach distant goals (mean terminal distance of 3.12. After adaptation, all methods recover accurate trajectories, with Last-Branch (0.45) slightly outperforming SetONet-FT (0.48), indicating that the pretrained basis functions (trunk network) generalize well and only the coefficients require updating. (b) Held-out obstacle avoidance tasks. All cost-based fine-tuning strategies reduce collision time well below the pretrained model
Figure 8: Out-of-distribution fine tuning on a Quadrotor task. Left two panels show multiple representative trajectories for SetONet-FT and SetONet-Meta-Full adapting to a shifted goal location. The bar plot compares relative error across all fine-tuning and meta-learning methods.
Figure 9: Control predictions on a held-out HalfCheetah-v3 task, showing the first three control dimensions ( , rows) for three methods (columns) over time. Each panel compares the method’s prediction (colored) against the expert (black) after 25 gradient steps using 10 expert trajectory demonstrations of 100 timesteps each. SetONet-FT most closely tracks the expert across all dimensions.
Figure 10: Adaptation performance on held-out HalfCheetah-v3 tasks as a function of the number of expert demonstrations (panels) and gradient steps (horizontal axis). Lines show mean relative error over 5 seeds; shaded regions indicate standard deviation. With few gradient steps (1–5), the meta-trained variants outperform SetONet-FT, particularly in the low-data regime (1 and 5 demos). As the adaptation budget increases, MAML and SetONet-Meta-Full improve steadily while SetONet-FT has the largest decrease in error as the number of available samples is increased.