Neural Operators for Multi-Task Control and Adaptation

2026-04-03

AI Summary

This work addresses the challenge of generalization and rapid adaptation in multi-task optimal control, where the goal is to map task descriptions directly to optimal feedback policies. The authors propose a permutation-invariant neural operator architecture that learns this mapping end-to-end via behavioral cloning. Built upon a branch-trunk network design, the approach supports flexible adaptation strategies ranging from lightweight parameter updates to full-network fine-tuning, and incorporates meta-learning-based initialization to enable efficient few-shot adaptation. Evaluated across diverse parametric optimal control and locomotion benchmarks, the model demonstrates strong generalization to unseen tasks, out-of-distribution scenarios, and varying observation conditions, significantly outperforming existing meta-learning baselines in few-shot adaptation performance.

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.

Problem

• multi-task control
• optimal control
• task adaptation
• neural operators
• generalization

Innovation

• neural operators
• multi-task control
• permutation-invariant architecture
• few-shot adaptation
• branch-trunk structure