💻 computer science

Learning to Adapt: Representation-Based Reinforcement Learning for Multi-Task Skill Transfer

This paper proposes RepMT-SAC, a multi-task reinforcement learning framework that leverages spectral MDP decomposition to separate task-agnostic dynamics from task-specific adjustments, thereby achieving superior zero-shot performance and rapid few-shot adaptation in quadcopter trajectory-following tasks compared to existing baselines.

Original authors: Aryan Naveen, Haitong Ma, Haldun Balim, Na Li

Published 2026-06-12

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Aryan Naveen, Haitong Ma, Haldun Balim, Na Li

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are teaching a drone to fly. In the old way of doing things (standard Reinforcement Learning), if you wanted the drone to learn how to fly a specific path, you'd have to teach it from scratch. If you then asked it to fly a slightly different path, you'd have to start over and teach it again. It's like hiring a new tutor for every single math problem you want to solve; the student never really learns the concept of math, they just memorize the answers to specific questions. This is slow, wasteful, and the drone gets confused when faced with a new situation.

This paper introduces a new method called RepMT-SAC that teaches the drone to learn the "grammar" of flying first, so it can instantly understand new sentences (tasks) without needing to relearn the alphabet.

Here is how it works, broken down into simple concepts:

1. The Core Idea: Separating the "Engine" from the "Destination"

Think of the drone's physics (how it moves, how wind affects it, how heavy it is) as the engine. Think of the specific path it needs to fly as the destination.

Old Way: The drone tries to learn the engine and the destination all at once. If the destination changes, the whole learning process gets messy.
RepMT-SAC Way: The system splits the learning into two distinct parts:
1. The Task-Agnostic Core (The Engine): This part learns the universal rules of how the drone moves. It doesn't care where the drone is going, only how it moves. This is like learning how to drive a car regardless of whether you are going to the grocery store or the beach.
2. The Task-Specific Adjustment (The Destination): This part is a small, flexible "knob" that tells the engine where to go. It's like a GPS coordinate.

By separating these, the drone learns the hard part (how to fly) once, and then just tweaks the "GPS knob" for every new path.

2. The Two-Phase Training Process

The paper describes a two-step training camp for the drone:

Phase 1: The "Upstream" Boot Camp (Learning the Basics)
The drone is trained on a set of basic flight paths (Source Tasks). During this time, it learns the "universal engine" (the shared dynamics) and how to read different "GPS coordinates" (task encodings).

Analogy: Imagine a pilot student flying a simulator. They practice takeoffs, landings, and turns in various weather conditions. They aren't memorizing specific routes; they are mastering the feel of the plane. The paper uses a mathematical trick called "spectral decomposition" to ensure the student learns the feel of the plane separately from the specific route.

Phase 2: The "Downstream" Quick Adaptation (The Real Test)
Once the pilot is trained, you give them a brand new, complex flight path they have never seen before (Out-of-Distribution task).

The Magic: Because the pilot already knows how to fly the plane perfectly, they don't need to relearn everything. They just need to adjust their "GPS knob" slightly to match the new route.
Result: The drone can adapt to these new, difficult paths with very little extra practice (only 10% of the original training time).

3. The Results: Why It's Better

The researchers tested this on a quadcopter (a small drone) trying to follow different 3D paths. They compared their method against standard AI methods (SAC) and other advanced methods (CTRL).

On Known Paths (In-Distribution): The new method was perfect. It achieved a 100% success rate, while the standard method only succeeded about 60% of the time and was much less stable.
On New, Weird Paths (Out-of-Distribution): When given a path it had never seen before, the new method adapted quickly and also reached a 100% success rate after a tiny bit of extra practice. The standard methods struggled, often failing completely or flying erratically.

4. The Takeaway

The paper claims that by mathematically separating "how the world works" (dynamics) from "what we want to achieve" (rewards/tasks), robots can learn much faster and generalize better.

Instead of memorizing a million different flight paths, the drone learns the structure of flight. This allows it to look at a completely new, complex path and say, "I know how to fly; I just need to adjust my target slightly," rather than saying, "I have no idea what to do."

In short: It turns a robot that memorizes answers into a robot that understands the subject, allowing it to ace any test, even the ones it hasn't seen before.

Technical Summary: RepMT-SAC for Multi-Task Skill Transfer

Problem Statement

Reinforcement learning (RL) has demonstrated success in learning complex control policies for single tasks but suffers from sample inefficiency and poor generalization when applied to multi-task settings. Standard approaches often treat tasks independently, leading to policies that overfit to specific reward signals and fail to exploit shared structures, such as common environment dynamics or physical laws. The core challenge is disentangling task-specific objectives from task-invariant structures to enable the learning of reusable representations. This paper addresses a specific class of multi-task Markov Decision Processes (MDPs) where all tasks share identical underlying system dynamics but differ in their reward functions or reference objectives, a scenario common in control applications like quadrotor trajectory following.

Methodology: RepMT-SAC

The authors propose RepMT-SAC (Representation-based Multi-Task Soft Actor-Critic), a framework designed to explicitly disentangle task-invariant environment representations from task-specific objectives. The method is grounded in the spectral decomposition of MDPs and operates in two distinct phases:

1. Theoretical Foundation

The framework extends spectral decomposition to the multi-task setting. It posits that for any task $\tau$ , the reward function $r(s, a, \tau)$ can be expressed as a linear combination of task-invariant features $\phi(s, a)$ and task-dependent parameters $\theta(\tau)$ . Consequently, the action-value function $Q^\pi(s, a; \tau)$ admits a linear parameterization:
$Q^\pi(s, a; \tau) = \langle \phi(s, a), w^\pi(\tau) \rangle$
Here, $\phi(s, a)$ captures the shared, task-agnostic dynamics of the environment, while $w^\pi(\tau)$ encodes the task-specific contributions. This structure allows the Q-function to be modeled as a linear combination of reusable skills weighted by the task.

2. Two-Phase Learning Paradigm

Upstream Phase (Joint Learning): The framework jointly optimizes task-agnostic representations of the environment dynamics and task-dependent encodings of the Q-function.
- Dynamics Learning: It employs a spectral conditional density estimation framework to learn the task-invariant feature map $\phi(s, a)$ and an auxiliary representation $\mu(s')$ by approximating the conditional transition density $P(s'|s, a)$ . This is achieved via a sampling-based cross-entropy loss.
- Value Learning: The task encoding $w(\tau; \theta)$ is updated using a temporal-difference (TD) loss, ensuring the policy evaluation remains a stable linear optimization problem.
- Policy Derivation: A maximum-entropy policy is derived from the learned Q-function, following the Soft Actor-Critic (SAC) framework to encourage exploration and robustness.
Downstream Phase (Fast Adaptation): When encountering new, unseen tasks (out-of-distribution), the learned task-invariant dynamics ( $\phi$ and $\mu$ ) are frozen. Only the task-specific encoding $w(\tau; \theta)$ and the policy $\pi$ are updated via fine-tuning. This significantly reduces the number of trainable parameters, enabling rapid specialization with minimal environment interactions.

Key Contributions

Explicit Disentanglement: Unlike prior methods that implicitly entangle dynamics and rewards within latent embeddings, RepMT-SAC explicitly parameterizes task-invariant dynamics as a standalone object, separating them from task-specific objectives.
Spectral MDP Decomposition: The paper extends spectral decomposition theory to multi-task RL, proving that task-conditioned value functions can be linearly parameterized by task-invariant features and task-dependent weights.
Efficient Adaptation Mechanism: By freezing the dynamics representation during downstream adaptation, the framework achieves rapid few-shot learning for out-of-distribution tasks without relearning the underlying physics of the system.
Application to Quadrotor Control: The method is applied to quadrotor trajectory following, where tasks are defined by different reference trajectories (encoded via Legendre polynomial coefficients) but share identical flight dynamics.

Experimental Results

The framework was evaluated in the IsaacSim environment on quadrotor trajectory-following tasks, comparing RepMT-SAC against standard SAC and an ablated variant (CTRL-SAC).

Source Tasks (In-Distribution Training): RepMT-SAC achieved a 100% success rate with a mean reward of $1094.2 \pm 17.9$ , significantly outperforming SAC ( $857.1 \pm 277.9$ ) and CTRL-SAC ( $377.9 \pm 35.4$ ). It demonstrated stable, monotonic learning with low variance compared to the oscillatory behavior of baselines.
In-Distribution (ID) Generalization: On tasks formed by convex combinations of source tasks, RepMT-SAC maintained a 100% success rate and a mean reward of $1100.7 \pm 16.1$ . In contrast, SAC achieved only 60% success, and CTRL-SAC failed entirely (0% success), highlighting the necessity of the structured decomposition for generalization.
Out-of-Distribution (OOD) Adaptation: For tasks defined by higher-order polynomials outside the training span, RepMT-SAC was fine-tuned for only 25,000 iterations (10% of the pretraining budget). It achieved a 100% success rate and a mean reward of $1034.8 \pm 23.3$ after fine-tuning. Baselines struggled to adapt, with SAC reaching only 40% success and CTRL-SAC remaining at 0%.
Performance Gain: The authors report that RepMT-SAC outperforms baselines by up to 30% in reward and success rate across various metrics.

Significance and Claims

The paper claims that RepMT-SAC provides a robust foundation for multi-task reinforcement learning by leveraging the shared structure of environment dynamics. Its primary significance lies in enabling strong zero-shot performance on in-distribution tasks and rapid few-shot adaptation to out-of-distribution tasks. By decoupling the learning of reusable dynamics from task-specific objectives, the framework overcomes the sample inefficiency and poor generalization typical of standard multi-task RL approaches.

The authors note that the current work relies on hand-crafted task distributions ( $\mu_{train}$ ) and future work will focus on leveraging the learned task embeddings to guide structured exploration and facilitate Sim2Real transfer to physical quadrotor platforms.