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 ci . Right: The corresponding optimal policies π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 Tθ[ℓi]→π^i , with pointwise evaluations of ℓi (in red) and of π^i at the point y . The branch network maps sensor locations (x,u) of a cost function ℓ(x,u;ϕ) to task-dependent coefficients {ck(ℓ)}k=1p . The red points indicate the pointwise samples of ℓi . The trunk maps query locations y=(x,t) to a set of learned basis functions {bk(y)}k=1p . Their inner product yields the predicted control output at the query location y
Figure 3: Overview of the meta-training procedure. The inner loop adapts the parameters θk to each task i using the support set Ditr . SetONet-Meta updates only the branch coefficients θbranch , while SetONet-Meta-Full updates all parameters θ . The adapted parameters θi′ from a batch of B tasks are evaluated on held-out query sets Dieval , and the outer loop computes the meta-gradient to update the shared initialization θk+1 .
Environment
dx/du
Expert
Tasks
Traj.
Task parameters
P2P-Cost
4/2
LQR
500
100
xg∈[−10,10]2
P2P-Cost-Small
4/2
LQR
50
10
xg∈[−10,10]2
P2P-Dynamics
4/2
iLQR
100
100
μ,vmax,amax
Quadrotor
6/2
iLQR
100
20
m,L,α
Obstacle
4/2
NLP
500
60
nobs,(x,y)i
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 L2 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 L2 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.
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 x -axis showing the method’s relative L2 error and the y -axis showing MAML’s error on the same task. Points above the diagonal ( y=x ) 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 L2 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 L2 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 ( u1,u2,u3 , 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 L2 error over 5 seeds; shaded regions indicate ±1 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.
We study the approximation and statistical complexity of learning collections of operators in a shared multi-task setting, with a focus on the Multiple Neural Operators (MNO) architecture. For broad classes of Lipschitz multiple operator maps, we derive near-optimal upper bounds for approximation and statistical generalization. On the lower-bound side, we establish a curse of parametric complexity and prove corresponding minimax rates. Together, these results show that shared representations across tasks do not increase the overall cost: multi-task operator learning follows the same scaling laws as single operator learning. We also compare MNO with a multi-task extension of DeepONet based on concatenated task inputs and show that, from a worst-case approximation-complexity perspective, both architectures satisfy essentially the same asymptotic rates.
Adrien Weihs, Hayden Schaeffer
Department of Mathematics, University of California Los Angeles, Los Angeles, CA 90095, USA.
Reinforcement learning has achieved remarkable success in learning complex control policies, yet its applicability remains limited due to sample inefficiency and poor generalization across tasks. In this work, we propose RepMT-SAC, a framework for multi-task RL that enables efficient knowledge sharing and robust transfer to new tasks. RepMT-SAC uses spectral MDP decomposition to capture transferable dynamics, structuring the value function into a task-agnostic core with a minimal task-specific adjustment. This design allows for strong zero-shot performance on in-distribution tasks and rapid few-shot adaptation to out-of-distribution tasks. We evaluate RepMT-SAC on quadcopter trajectory-following tasks across in-distribution and out-of-distribution contexts, demonstrating that it outperforms baselines by up to 30%.
Aryan Naveen, Haitong Ma, Haldun Balim +1
Massachusetts Institute of Technology · Harvard School of Engineering and Applied Sciences
Many physical AI tasks require sequential implicit computation: at each step, boundary controls are applied, and the resulting configuration is obtained by solving an equilibrium problem. This setting arises naturally in deformable object manipulation, where even bending a deformable linear object (DLO) to a target shape can be nonlinear and multistable: identical boundary conditions may produce different configurations depending on actuation history. Unlike explicit transition models, the control-to-configuration relation is implicit and history-dependent, making long-horizon learning and control brittle; backpropagating through iterative solves is also memory- and compute-intensive. We propose Neural Control, a boundary-control framework that propagates gradients through branch-dependent sequences of equilibrium solves rather than a single fixed point. Neural Control computes trajectory-dependent proxy gradients by differentiating equilibrium conditions with an adjoint formulation, avoiding solver unrolling while keeping forward rollouts on converged equilibria. Combined with receding-horizon continuation, Neural Control re-anchors optimization to realized equilibria and mitigates basin switching. We validate Neural Control on simulated and real DLO manipulation, compare against SPSA and iCEM, and demonstrate applicability to a learned DEQ-style implicit equilibrium model.
Dezhong Tong, Jiawen Wang, Hengyi Zhou +3
The Robotics Department, University of Michigan, Ann Arbor, US · The Department of Mechanical and Aerospace Engineering, University of California, Los Angeles, US