Multi-Task RL — Annotated Lecture Notes & Key Takeaways

This lecture wraps up the course’s unit on model-based reinforcement learning and dives into multi-task RL, the central topic for the rest of the session. The lecturer walks through how learned world models can generate synthetic training data to cut down on real-world interaction, then builds up to training generalist policies that handle many tasks at once. You’ll also get a breakdown of hindsight relabeling, a data-reuse trick that the course will have you implement for an assignment.
 

Wrapping Up Model-Based Reinforcement Learning

Last lecture covered planning with learned models: using a predictive world model to simulate future outcomes and pick actions on the fly. That approach is simple and flexible for changing goals, but computationally heavy at test time and only practical for short time horizons.

This lecture focuses on the second major use case for learned models: generating synthetic data to train your policy. The core idea is that you don’t have to learn your policy only on real environment data. You can use your learned simulator to imagine extra transitions and add those to your training set.

The standard algorithm follows a loop. First, collect a dataset of real trajectories from the environment. Next, update your world model to match that real data as closely as possible. Then, sample random states from your real dataset, and roll out short trajectories starting from those states entirely inside the learned model. Finally, combine both the real data and the generated synthetic data to update your policy, then repeat the whole cycle.

The lecturer emphasized keeping these synthetic rollouts very short. World models get less accurate the further you project into the future, so long imaginary trajectories drift away from real dynamics and can mislead your policy. Short rollouts just add a little extra state coverage around your real data points, which tends to work much better in practice.

There are big upsides to this approach. Most importantly, you need far less real environment data. That’s a massive deal if collecting real data is dangerous, expensive, or slow — like running experiments on physical robots. You can even train the world model on data that has no reward labels at all, then use a small set of rewarded data for the policy head. That lets you leverage huge unlabeled interaction datasets that standard model-free RL couldn’t use. World models also tend to be somewhat task-agnostic, so you can often transfer a well-trained model across different reward functions or tasks.

There are significant downsides too. The model is trained to predict the world accurately, not to produce the best policy learning signal. That means it can waste capacity on irrelevant details — like matching exact pixel colors in an image observation — that don’t actually help you make better decisions. In some domains, learning an accurate world model is actually harder than just learning the policy directly. The lecturer used the example of pouring water: the policy is simple (just rotate your wrist) but modeling fluid dynamics accurately is extremely complex. You also pick up a lot of extra hyperparameters: model architecture, training length, rollout horizon, how much synthetic data to generate. Tune those poorly, and your performance can get worse than just using real data alone.

In robotics, the track record lines up with these tradeoffs. Model-based RL has been very successful for legged locomotion and dexterous manipulation tasks where you have compact state representations — joint angles, object positions, that sort of thing. It’s been much less successful with high-dimensional image observations, though the lecturer noted that rapid improvements in generative video models might shift that in the future.

The lecturer also briefly touched on other types of models you can learn besides forward prediction: inverse models (given two consecutive states, what action caused the transition), multi-step inverse models, and generative transition models. These all train on plain transition data without rewards, and can be useful for different kinds of control setups.
 

What Multi-Task RL Actually Means

Multi-task reinforcement learning is all about training one generalist policy that can perform many different tasks, instead of building a separate policy for every single thing you want an agent to do.

The lecturer gave a range of examples. A legged robot that can walk, run, crouch, and dance. A mobile robot that can hang towels, unload dishwashers, and wipe up spills. A language assistant that can book travel, order groceries, and answer questions. A recommendation system that personalizes results for many different users. Even an agent that can play dozens of different video games.

There are two core motivations for doing this. First, data efficiency. Most tasks share underlying structure. Walking forward and walking backward both rely on the same core balance skills. Learning to interact with one website teaches you patterns that transfer to other sites. Training all tasks together lets you amortize that shared learning, so you need far less data per individual task. Second, empirical results keep showing that generalist systems often end up more reliable and higher-performing than single-task policies trained on the same data. The lecturer called this one of the cooler findings from recent AI research, and a big extra incentive to move beyond single-task training.

Formally, each task is defined as its own Markov Decision Process (MDP). That means each can have its own state space, action space, initial state distribution, transition dynamics, and reward function. They don’t have to all be the same. For example, controlling two different robot arms would mean different action dimensions and different dynamics. Walking on grass vs. flat concrete has the same reward function but different transition dynamics. Different users for a recommendation system have completely different reward functions.

To make this work, you need some way to tell the policy which task it’s supposed to be doing at any given time — a task identifier. The simplest version is just an index number assigned to each task. More expressive options include natural language descriptions, video demonstrations, or a target goal state. When you use a target goal state as the task identifier, that’s called goal-conditioned RL, and reward is typically defined as negative distance between the current state and the goal.

One way to frame this is that you can just append the task identifier to the state, turning all your separate task MDPs into one big aggregated MDP. In theory, you could then throw any standard RL algorithm at it and call it a day. But the lecturer made clear that you can get much better results by actually leveraging the shared structure between tasks.
 

Multi-Task Imitation Learning Basics

If you’re working with imitation learning (behavior cloning) instead of full RL, the multi-task extension is conceptually very simple. You have demonstration datasets for each of your tasks. You train a single policy that takes both the state and the task identifier as input, and minimize the average action prediction loss across all tasks.

In practice, there are a couple common pitfalls. If your tasks demand very different behaviors — say, one tells the robot to move forward and another tells it to lie down — you can get really high gradient variance during training. A minibatch full of forward-moving examples will push the weights one way, then a minibatch full of lying-down examples will yank them the other way. The policy can theoretically use the task ID to tell the two apart, but it can take a long time for the network to learn that conditioning properly.

The most useful fix for this is stratified minibatching. Make sure every single training batch contains roughly equal amounts of data from each task. That evens out the gradient updates and keeps training stable. The lecturer also recommended normalizing states and actions per task, so they all sit on similar numerical scales.

Architecture design gets interesting when task identifiers use different modalities. For example, your state might be camera images, and your task descriptor might be a text command or a short video of a human doing the task. The lecturer walked through an older robot manipulation paper that used separate encoders for text and video task inputs, then combined those embeddings with the image state backbone. More modern approaches generally use transformer architectures, treating the task description like a prompt — very similar to how instruction-tuned language models work.

There’s a natural tradeoff between one shared policy and separate policies per task. Separate policies use a lot more memory, switch slowly between tasks, and can’t share data between tasks, so they almost always perform worse when data is limited. But cramming too many different tasks into one shared model can create optimization conflicts and hurt performance on individual tasks. In some cases, especially when tasks are extremely different, separate policies can still be the right call.
 

Hindsight Relabeling: Reusing Data Across Tasks

When you move from imitation learning to full reinforcement learning, you run into a question about data sharing. If you use separate replay buffers for each task, you’re wasting a lot of experience that could be useful across tasks.

The lecturer used a hockey example to make this concrete. Imagine you’re training an agent to shoot on goal. On one attempt, the agent messes up the shot, but accidentally sends a perfect pass to a teammate. That trajectory is low-reward for the shooting task, but it’s actually a great example of good passing. If you just leave it in the shooting buffer, you’re throwing away useful training signal for the passing task.

Hindsight relabeling is the trick to fix this. You take that trajectory, recalculate what the reward would be for the passing task, swap out the task identifier, and add it to the passing task’s replay buffer. The actual state and action transitions stay exactly the same — you’re just reinterpreting what the trajectory means for a different task.

The full training loop works like this. Pick a task, collect trajectories using your current policy conditioned on that task, and store them in the replay buffer. Then, relabel some of those trajectories for other tasks by recomputing their rewards under the new task’s reward function. Add those relabeled trajectories to the appropriate buffers, then update your policy on all the accumulated data, and repeat.

This approach only works under specific conditions, which the lecturer walked through. First, tasks have to share the same state and action space. If they have different dimensions, you can’t just drop a trajectory from one into the other’s buffer. Second, the transition dynamics have to be identical across tasks. You’re only relabeling rewards, not the physics of how actions change the state. If dynamics differ between tasks, the relabeled trajectory won’t be valid. Third, reward functions need to be different — that’s actually when this technique adds the most value. If rewards were the same, there’d be no point in relabeling.
Finally, you have to use an off-policy reinforcement learning algorithm. The trajectory was collected by a policy trying to do task i, not task j. That means it’s off-policy data for task j, so on-policy algorithms won’t handle it correctly. I found this point a little tricky to wrap my head around at first, but the hockey example clarifies it: the actions you take when you’re trying to shoot aren’t the same actions you’d choose if you were intentionally trying to pass. The data doesn’t reflect what the task-j policy would do.

One open, unresolved point: there’s no great principled way to decide which tasks to relabel to, or how often to do it. Right now it’s mostly heuristic. You can relabel randomly, or prioritize trajectories that happen to get high reward on the new task. There’s no widely agreed optimal strategy yet, which makes it an active research area.

Goal-conditioned RL is a particularly natural fit for hindsight relabeling. When reward is just negative distance to a goal state, you can relabel any trajectory with any goal state almost for free. This is the core idea behind hindsight experience replay, a very popular algorithm for goal-reaching tasks.

Multi-task RL is one of the most active areas in modern reinforcement learning, because it pushes us toward more general, capable agents that don’t have to be trained from scratch for every new problem. The tools from this lecture — synthetic data from world models, stratified training for multi-task policies, and hindsight relabeling — all build toward that goal of making RL more data efficient and more general.