Lecture 12: Multi-Task & Goal-Conditioned RL

From Single-Task to Multi-Task RL

Throughout the first part of this course, we trained agents to solve one task at a time: balance a cartpole, navigate a maze, or achieve a single goal configuration. But the real world demands versatility. A household robot must make coffee, clean dishes, and set the table—all with the same body and brain. A warehouse system must pick, place, pack, and sort thousands of different products. Training a separate policy for each task from scratch is both data-inefficient and fails to leverage the obvious structural similarities between related tasks.

Multi-task reinforcement learning addresses this by training a single policy (or a small family of policies) to solve an entire distribution of tasks simultaneously. The central promise is twofold: (1) sharing data and representations across tasks can dramatically improve sample efficiency, and (2) the resulting generalist agent can exhibit combinatorial generalization—solving novel task combinations never seen during training.

Key Insight

The shift from single-task to multi-task RL mirrors the transition from training one neural network per ImageNet class to training a single classifier for all classes. Just as shared low-level features (edges, textures) benefit all vision tasks, shared dynamics understanding and motor primitives benefit all control tasks.

Why Multi-Task RL?

Several compelling motivations drive multi-task RL research:

Data efficiency through sharing: Experience collected while learning one task can accelerate learning on related tasks. A robot learning to grasp mugs generates trajectories that are also informative for grasping bowls.
Amortization of computation: Rather than running a full RL optimization for each new task, a multi-task agent can reuse learned representations and adapt quickly.
Generalization to new tasks: With sufficient task diversity during training, the agent may generalize to tasks it has never encountered before, provided they lie within the support of the training distribution.
Practical deployment: Deploying a single generalist model is far simpler than maintaining a library of specialized policies.

Formalizing Multi-Task RL

To formalize multi-task RL, we need to specify what constitutes a "task" and how tasks relate to one another. The standard approach is to define each task as a distinct Markov Decision Process drawn from a shared family.

Definition

Task Distribution. A task distribution $p(\mathcal{M})$ is a distribution over MDPs $\mathcal{M}_i = (\mathcal{S}, \mathcal{A}, T_i, R_i, \gamma)$, where different tasks may differ in their transition dynamics $T_i$, their reward function $R_i$, or both, while sharing the same state space $\mathcal{S}$ and action space $\mathcal{A}$.

Each task $\mathcal{M}_i$ is associated with a task identifier $z_i$ that encodes the relevant information distinguishing it from other tasks. This identifier might be a one-hot vector, a natural language description, a goal image, or a learned latent embedding.

The Task-Augmented MDP

A clean way to incorporate the task identifier into the standard MDP framework is through state augmentation. We define an augmented state that concatenates the original environment state $\bar{s}$ with the task identifier $z_i$:

$$s = (\bar{s}, z_i)$$

Under this formulation, a single MDP $\mathcal{M}^{\text{aug}} = (\mathcal{S} \times \mathcal{Z}, \mathcal{A}, T, R, \gamma)$ subsumes the entire task family. The transition dynamics and reward function of the augmented MDP dispatch to the appropriate task-specific behavior based on $z_i$. This is elegant because it reduces multi-task RL to single-task RL in a larger state space—at least in principle. In practice, the challenge lies in learning a policy that can effectively condition on $z_i$ to exhibit diverse behaviors.

What Varies Across Tasks?

Different problem settings emphasize different sources of variation:

Different rewards, same dynamics: The environment physics are identical, but the agent is asked to achieve different objectives. Example: a robot arm with fixed dynamics but varying target object placements.
Different dynamics, same reward: The objective is the same, but the environment behaves differently. Example: locomotion across varying terrain types, all with the goal of moving forward.
Both vary: The most general case, where each task has unique dynamics and rewards. Example: different cooking recipes requiring different ingredients and sequences of actions.

Example

Robot Manipulation Tasks. Consider a robotic arm that must perform various manipulation tasks: pick up a red block, push a blue cylinder to a target location, or stack objects in a specified order. Here $\bar{s}$ includes joint angles and object poses, while $z_i$ encodes the specific task instruction. All tasks share the same robot dynamics, but the reward function (what constitutes success) differs for each task.

Goal-Conditioned Reinforcement Learning

An important special case of multi-task RL arises when the task identifier is simply a desired goal state $s_g \in \mathcal{S}$. This is called goal-conditioned RL, and it has become one of the most widely studied formulations in the field.

Definition

Goal-Conditioned Policy. A goal-conditioned policy $\pi(a \mid s, g)$ maps the current state $s$ and a desired goal $g$ to an action distribution. The reward function is typically defined as a function of the distance between the achieved state and the goal: $$r(s, a, g) = -d(s', g)$$ where $d$ is a distance metric (e.g., Euclidean distance) and $s'$ is the resulting next state. In the sparse variant, $r = \mathbf{1}[d(s', g) < \epsilon]$ for some threshold $\epsilon$.

Goal-conditioned RL is appealing because goals provide a natural, interpretable way to specify tasks, and the space of goals is often the same as the state space, enabling a form of combinatorial generalization. A policy trained to reach a diverse set of goal positions can interpolate to reach goals it never explicitly practiced during training.

Universal Value Functions

Schaul et al. (2015) introduced the concept of Universal Value Function Approximators (UVFAs), which extend the standard value function to be conditioned on both state and goal:

$$V(s, g; \theta) = \E\!\left[\sum_{t=0}^{\infty} \gamma^t r(s_t, a_t, g) \;\middle|\; s_0 = s, \pi(\cdot \mid \cdot, g)\right]$$

Similarly, the action-value function becomes $Q(s, a, g; \theta)$. The key insight is that a single neural network can learn to represent value functions across the entire goal space, with the network implicitly learning to share representations for similar goals. This is precisely the multi-task state augmentation idea, specialized to goals.

Multi-Task Imitation Learning

Before diving into multi-task RL algorithms, it is worth noting that multi-task imitation learning—learning from demonstrations rather than from reward signals—provides a simpler starting point and has seen enormous practical success in robotics.

Given a dataset of expert demonstrations across multiple tasks, $\mathcal{D} = \{(\tau_j, z_j)\}_{j=1}^{N}$ where $\tau_j$ is a trajectory and $z_j$ is the corresponding task identifier, we train a task-conditioned policy via behavioral cloning:

$$\max_\theta \; \E_{(\tau, z) \sim \mathcal{D}} \left[\sum_{t} \log \pi_\theta(a_t \mid s_t, z)\right]$$

Stratified Sampling

A practical challenge in multi-task imitation learning is data imbalance. Some tasks may have far more demonstration data than others. Naively sampling uniformly from $\mathcal{D}$ will bias the policy toward well-represented tasks. Stratified sampling addresses this by ensuring each task receives equal representation in each training batch, regardless of how many demonstrations are available for it.

Key Insight

Stratified sampling prevents the multi-task policy from "forgetting" rare tasks. Without it, tasks with abundant data dominate the gradient, and the policy effectively becomes a specialist on the most common tasks. This is analogous to class-balanced sampling in long-tailed image recognition.

Task Conditioning Architectures

How should the task identifier $z_i$ be injected into the policy network? Several architectures have proven effective:

Concatenation: The simplest approach appends $z_i$ to the input alongside the state $s$. The network processes $[s; z_i]$ through shared layers.
FiLM conditioning: Feature-wise Linear Modulation layers allow $z_i$ to scale and shift intermediate activations: $\text{FiLM}(h; z) = \gamma(z) \odot h + \beta(z)$, where $\gamma$ and $\beta$ are learned functions of $z$.
Language conditioning: When tasks are described in natural language, a language encoder (e.g., a pretrained transformer) maps the instruction to an embedding that conditions the policy. This enables zero-shot generalization to novel task descriptions.

Example

BC-Z (Jang et al., 2022). BC-Z trains a single vision-language-conditioned policy on over 100 manipulation tasks. The task identifier is a natural language instruction such as "pick up the red block" or "open the drawer." At test time, the robot can follow instructions for tasks never seen during training, demonstrating zero-shot generalization from language-conditioned imitation learning.

More recently, OpenVLA and similar vision-language-action models take this further by using large pretrained vision-language models as the backbone, with robot actions generated as output tokens. These models leverage internet-scale pretraining to achieve broad generalization across manipulation tasks.

Multi-Task RL Algorithms

Extending RL algorithms to the multi-task setting is conceptually straightforward: we condition both the policy and the value function on the task identifier $z_i$, and we optimize across the task distribution. Concretely, the multi-task RL objective is:

$$\max_\theta \; \E_{\mathcal{M}_i \sim p(\mathcal{M})} \left[\E_{\pi_\theta(\cdot \mid \cdot, z_i)} \left[\sum_{t=0}^{\infty} \gamma^t r_i(s_t, a_t)\right]\right]$$

The policy $\pi_\theta(a \mid s, z_i)$ and Q-function $Q_\theta(s, a, z_i)$ are both conditioned on the task. In practice, we can apply any standard RL algorithm—SAC, PPO, DQN—with task-conditioned networks. During training, we sample tasks from $p(\mathcal{M})$, collect experience, and update the shared parameters $\theta$.

Challenges in Multi-Task RL

While the formulation is clean, several practical challenges arise:

Conflicting gradients: Different tasks may push the shared parameters in opposing directions, leading to destructive interference and suboptimal performance on some tasks.
Reward scale variation: Tasks with higher-magnitude rewards dominate the optimization. Reward normalization per task is essential.
Exploration complexity: The agent must explore effectively across all tasks simultaneously. A strategy that is informative for one task may be useless for another.
Negative transfer: In some cases, training on additional tasks can actually hurt performance compared to single-task training, particularly when tasks are dissimilar.

Hindsight Experience Replay

One of the most elegant ideas in goal-conditioned RL is Hindsight Experience Replay (HER), introduced by Andrychowicz et al. (2017). HER addresses a critical problem: when goals are specified with sparse binary rewards ($r = 1$ if the goal is reached, $0$ otherwise), the agent almost never receives positive reward during early training, making learning extremely difficult.

The key insight is deceptively simple: even if the agent fails to reach its intended goal, the trajectory it produced did reach some state. We can retroactively relabel that trajectory as if the achieved state had been the goal all along.

Definition

Hindsight Experience Replay (HER). Given a trajectory $\tau = (s_0, a_0, s_1, a_1, \ldots, s_T)$ collected while attempting to reach goal $g$, HER creates additional training examples by substituting alternative goals $g' \in G(\tau)$, where $G(\tau)$ is a set of goals derived from the trajectory itself. The most common strategy sets $g' = s_T$ (the final state actually reached), yielding a "successful" trajectory with reward signal.

Hindsight Experience Replay (HER)

Sample a goal $g$ from the goal distribution.
Collect a trajectory $\tau = (s_0, a_0, s_1, \ldots, s_T)$ using the current goal-conditioned policy $\pi(a \mid s, g)$.
Store the original transitions $\{(s_t, a_t, r_t, s_{t+1}, g)\}$ in the replay buffer with the original goal $g$ and rewards $r_t = r(s_t, a_t, g)$.
For each transition, also store relabeled versions $\{(s_t, a_t, r'_t, s_{t+1}, g')\}$ with alternative goals $g' \in G(\tau)$ and recomputed rewards $r'_t = r(s_t, a_t, g')$.
Train the goal-conditioned Q-function and policy using the augmented replay buffer.

Common relabeling strategies for $G(\tau)$: final (use $s_T$), future (sample a state reached later in the same trajectory), episode (sample any state from the trajectory), or random (sample from all collected states).

Key Insight

HER transforms every trajectory into a learning signal, regardless of whether the original goal was achieved. This is particularly powerful in sparse-reward settings where standard RL algorithms would receive no gradient signal for thousands of episodes. The "future" relabeling strategy is generally the most effective, as it creates a natural curriculum from easy (nearby) to hard (distant) goals.

HER for Multi-Task RL

The hindsight relabeling idea extends naturally beyond goals to general multi-task settings. If the task identifier $z_i$ determines the reward function $r_i$, and we can compute the reward under alternative task identifiers, then we can relabel any trajectory with alternative tasks:

$$\text{Original: } (s_t, a_t, r_i(s_t, a_t), s_{t+1}, z_i) \quad \longrightarrow \quad \text{Relabeled: } (s_t, a_t, r_j(s_t, a_t), s_{t+1}, z_j)$$

This multi-task hindsight relabeling effectively multiplies the useful data by the number of tasks, since each trajectory provides learning signal for many different task identifiers. The only requirement is that we can efficiently evaluate the reward function for alternative tasks given the observed transitions.

Practical Architectures and Scaling

Scaling multi-task and goal-conditioned RL to complex environments with high-dimensional observations (e.g., images) and large task spaces requires careful architectural choices.

Vision-Language-Action Models

The most recent wave of multi-task robot learning leverages large pretrained models. The recipe combines three ingredients:

Visual encoder: A pretrained vision model (e.g., ViT, SigLIP) that maps raw images to a compact representation.
Language encoder: A pretrained language model that embeds task instructions.
Action decoder: A network that maps the fused vision-language representation to robot actions, often trained with a flow-matching or diffusion objective to capture multimodal action distributions.

Systems like $\pi_0$ (Physical Intelligence, 2024) demonstrate that this approach can yield a single policy capable of folding laundry, packing boxes, busing tables, and assembling objects—all from a unified architecture trained on diverse multi-task data. The scale of the training data (thousands of hours of robot teleoperation across many tasks) is critical for the generalization to emerge.

Challenges at Scale

As the number of tasks grows, new challenges emerge beyond those discussed earlier:

Catastrophic forgetting: Revisiting old tasks less frequently can cause the policy to degrade on them. Replay buffers and careful curriculum design help mitigate this.
Task specification ambiguity: Natural language instructions can be ambiguous. "Clean up" might mean different things in different contexts. Grounding language in the physical world remains an open challenge.
Evaluation complexity: Measuring the performance of a policy across hundreds or thousands of tasks requires careful benchmarking and aggregation of metrics. Average success rate can mask catastrophic failure on specific tasks.

Theoretical Perspective

Multi-task RL has strong theoretical motivation. Under the task-augmented MDP formulation, we can analyze the sample complexity benefits of sharing data across tasks.

Consider $K$ tasks, each requiring $N$ samples to learn in isolation. If the tasks share a common representation $\phi: \mathcal{S} \to \R^d$ (where $d$ is much smaller than $|\mathcal{S}|$), then the total sample complexity for learning all $K$ tasks jointly can scale as $O(d \cdot K + d^2)$ rather than $O(N \cdot K)$. The first term accounts for task-specific learning on top of the shared representation, and the second for learning the representation itself. When $d^2 \ll N \cdot K$, multi-task learning provides a substantial sample complexity reduction.

Theorem

Benefits of Representation Sharing (Informal). If $K$ tasks share a common low-dimensional representation of dimension $d$, the total sample complexity for learning all tasks jointly is $\tilde{O}(dK + d^2)$, compared to $\tilde{O}(NK)$ for learning each task independently, where $N$ is the sample complexity per task. This yields a factor-of-$N/d$ improvement when $K$ is large.

This result formalizes the intuition that shared structure across tasks enables faster learning. The key assumption—that a shared low-dimensional representation exists—is both the strength and the limitation of the theory. In practice, the degree of sharing varies, and negative transfer can occur when tasks are insufficiently related.

Looking Ahead

This lecture introduced multi-task and goal-conditioned RL as natural extensions of the single-task RL framework. The task-augmented MDP formulation provides a unified view, and algorithms like HER demonstrate how to squeeze maximum learning signal from every trajectory.

Key takeaways:

Multi-task RL trains a single policy to solve a distribution of tasks by conditioning on task identifiers.
Goal-conditioned RL is an important special case where the task is specified by a desired state.
Hindsight relabeling (HER) enables learning from failure by retroactively reassigning goals or tasks to collected trajectories.
Modern architectures use vision-language-action models pretrained on internet-scale data to achieve broad generalization across manipulation tasks.

A natural question emerges: can we go beyond training on a fixed task distribution and instead learn an agent that rapidly adapts to entirely new tasks from just a few interactions? This is the promise of meta-reinforcement learning, which we take up in the next lecture.