🤖 machine learning

TSN-Affinity: Similarity-Driven Parameter Reuse for Continual Offline Reinforcement Learning

This paper proposes TSN-Affinity, a novel continual offline reinforcement learning method that leverages TinySubNetworks and a Decision Transformer to enable task-specific parameterization and similarity-driven knowledge sharing, offering a memory-efficient alternative to replay-based strategies that effectively mitigates catastrophic forgetting across discrete and continuous control tasks.

Original authors: Dominik Żurek, Kamil Faber, Marcin Pietron, Paweł Gajewski, Roberto Corizzo

Published 2026-04-29

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Dominik Żurek, Kamil Faber, Marcin Pietron, Paweł Gajewski, Roberto Corizzo

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 trying to teach a robot to play video games and move a robotic arm. The catch? You can't let the robot practice in the real world because it's too expensive or dangerous. Instead, you have to teach it using a library of old video recordings (datasets) of humans doing these tasks.

Now, imagine you have to teach this robot a new game every week. The big problem in "Continual Offline Reinforcement Learning" (CORL) is forgetting. If you teach the robot to play Breakout, it might get so good at that that it forgets how to play Pong. If you try to teach it Pong without deleting the old data, the robot gets confused and messes up both games.

This paper introduces a new method called TSN-Affinity to solve this problem. Here is how it works, using simple analogies:

1. The Old Way: The "One Big Brain" Problem

Most previous methods tried to teach the robot using one giant, shared brain. To prevent forgetting, they would save little snippets of old videos (replay) and mix them with new ones.

The Flaw: It's like trying to learn French while speaking English, then learning Spanish while speaking both. Eventually, the languages get mixed up, or you need a massive library of old tapes to keep everything straight, which takes up too much space.

2. The New Way: TSN-Affinity (The "Modular Toolbox")

The authors propose a smarter approach using TinySubNetworks. Imagine the robot's brain isn't one giant lump, but a toolbox with many small, specialized drawers.

Specialized Drawers (Subnetworks): When the robot learns a new task (like Breakout), it doesn't rewrite its whole brain. Instead, it opens a specific, tiny drawer and fills it with the exact tools needed for that game. Once that drawer is closed, it stays exactly as it is. If you learn a new game later, you open a different drawer. This means the robot never forgets the old games because the old tools are locked away and untouched.

3. The Secret Sauce: "Affinity Routing" (The Smart Librarian)

If every new task got its own brand-new drawer, the toolbox would get huge and messy. The paper's innovation is a Smart Librarian (called the Routing Mechanism) that decides which drawer to use.

Before the robot starts learning a new task, the Librarian asks: "Does this new task look like any of the old tasks we already have tools for?"

The Librarian checks two things:

Action Affinity (The Moves): "Do the moves required for this new game look like the moves we already know?" (e.g., If the new game requires jumping, and we already have a "Jumping" drawer, maybe we can use that).
Latent Affinity (The Feel): "Does the new task 'feel' similar to the old ones inside the robot's mind?" (e.g., Do the patterns of the game look similar, even if the graphics are different?)

The Decision:

If they are similar: The Librarian says, "Great! Let's reuse the existing drawer." The robot uses the old tools (which are frozen and safe) and just adds a few new, tiny tools on top. This saves space and improves performance.
If they are different: The Librarian says, "No, this is too different." It opens a brand-new, empty drawer for the new task.

4. The Results: Video Games vs. Robot Arms

The authors tested this on two very different challenges:

The Video Games (Atari): These are fast, pixel-based games with simple buttons (discrete actions).
- Result: The method was a huge success. The "Specialized Drawers" approach completely stopped the robot from forgetting old games. The "Smart Librarian" helped the robot perform even better by reusing the right tools, often matching the performance of a robot trained on just one game at a time.
The Robot Arm (Panda): This involves moving a physical arm in a 3D space with smooth, continuous movements (like picking up a cup).
- Result: This was much harder. The "Smart Librarian" had a harder time deciding which tools to reuse because the movements were so complex and varied. While the method still worked better than the old "One Big Brain" methods, it showed that reusing tools in complex physical tasks is tricky. The robot had to balance keeping old skills safe while learning new, difficult physical movements.

Summary

TSN-Affinity is like giving a student a set of separate, labeled notebooks instead of one giant textbook.

When they learn a new subject, they open a fresh notebook.
If the new subject is similar to an old one, they check if they can use the old notes as a base (reusing them) rather than starting from scratch.
This ensures they never forget what they learned in the past, and they don't get confused by mixing up different subjects.

The paper proves that for learning from old data without messing up the past, having a system that knows when to reuse old knowledge and when to start fresh is much better than just trying to memorize everything in one big pile.

1. Problem Statement

The paper addresses Continual Offline Reinforcement Learning (CORL), a setting where an agent must learn a sequence of tasks from static, pre-collected datasets without interacting with the environment. This scenario arises in domains where live interaction is too expensive, risky, or impossible (e.g., robotics, healthcare).

CORL faces a "dual difficulty":

Offline RL Challenges: Dealing with distributional shifts between training data and evaluation, and avoiding Out-of-Distribution (OOD) actions.
Continual Learning Challenges: Preventing catastrophic forgetting of previously learned tasks while acquiring new ones, without access to past data unless explicitly stored.

Existing solutions rely heavily on replay-based methods (storing and reusing past samples), which incur memory overhead and suffer from distribution mismatches. Conversely, architectural methods (which allocate specific parameters for tasks) have shown promise in supervised learning but remain underexplored in CORL due to the complexities of sequential decision-making and policy-induced distribution shifts.

2. Methodology: TSN-Affinity

The authors propose TSN-Affinity, a novel architectural approach that combines TinySubNetworks (TSN) with a Decision Transformer (DT) backbone. The core innovation is a similarity-driven parameter reuse strategy that routes tasks to specific model copies based on "affinity."

Core Components:

Decision Transformer Backbone: The method uses a DT to reformulate RL as a supervised sequence modeling problem. This naturally constrains predictions to the empirical distribution of the dataset, mitigating OOD issues.
- Atari (Discrete): Trained with Cross-Entropy loss on action logits.
- Panda (Continuous): Trained with Masked Mean-Squared Error on continuous actions.
Sparse Task-Specific Subnetworks: Instead of training a single dense model, the system maintains multiple model copies. For each task, a sparse binary mask is learned over the weights of a specific copy.
- Weight Protection: Once a weight is allocated to a task, it is "frozen" and protected from gradients during the training of subsequent tasks.
- Reuse: New tasks can "reactivate" frozen weights from previous tasks if they are deemed similar, using them as fixed functional components rather than retraining them.
Affinity Routing (The Novelty): Before training a new task $T_t$ $T_{t}$ , the system decides whether to reuse an existing model copy or spawn a new one. This decision is based on three routing modes:
1. Action Affinity ( $d_A$ ): Measures how well a source task's policy (subnetwork) predicts the actions of the new task's trajectories (using Cross-Entropy for Atari, MSE for Panda).
2. Latent Affinity ( $d_L$ ): Measures the similarity of the internal representations. It compares the latent statistics (mean and variance of observation encodings) of the new task against the stored statistics of source tasks using Symmetric KL Divergence.
3. Hybrid Affinity ( $d_H$ ): A weighted combination of Action and Latent affinities.

If the affinity score between the new task and a source task is below a threshold $\tau$ , the new task is assigned to that source's copy. Otherwise, a new copy is created.

Variants Studied:

TSN-Affinity: The full method with affinity-based routing.
TSN-ReplayKL: Replaces affinity routing with a heuristic based on observation-level similarity (KL divergence between replay memories).
TSN-Core: Removes routing entirely; uses a single evolving model copy with sparse subnetworks (no copy selection).

3. Key Contributions

TSN-Affinity Framework: A new CORL method that formulates knowledge transfer as a routing problem over sparse, task-specific subnetworks, dynamically selecting and reusing parameters based on RL-aware similarity.
Affinity Routing Mechanism: A specific mechanism that leverages action-level (policy compatibility) and representation-level (latent similarity) metrics to decide parameter reuse, moving beyond simple replay or random allocation.
Comprehensive Benchmarks: The introduction of two continual offline RL benchmarks covering:
- Visual Discrete Control: A suite of Atari games (Breakout, Alien, Atlantis, Boxing, Centipede).
- Continuous Robotic Manipulation: Panda-gym environments (Reach, Push, Pick-and-Place) using a Franka Emika Panda arm.
Empirical Validation: Extensive experiments demonstrating that architectural approaches can completely mitigate forgetting and that similarity-driven reuse significantly outperforms replay-based baselines, particularly in heterogeneous continuous-control settings.

4. Results

Atari Benchmark (Discrete Control)

Forgetting: Sparse subnetworks (TSN-Core) achieved zero forgetting across all tasks. Dense baselines (Naive, EWC, SI) suffered severe catastrophic forgetting (normalized returns ~15-18%).
Performance: While sparse subnetworks alone eliminated forgetting, Affinity Routing further improved multi-task performance.
- TSN-Affinity-A (Action affinity) and TSN-Affinity-L (Latent affinity) achieved ~89.7% and ~89.6% normalized performance, respectively.
- This outperformed the best dense baseline (Cumulative Replay at 76.2%) and approached the performance of single-task trained models.
Conclusion: In discrete settings, sparse subnetworks are sufficient to prevent forgetting, and affinity routing optimizes capacity usage to boost performance.

Panda Benchmark (Continuous Control)

Challenge: This setting is significantly harder due to heterogeneous action spaces and continuous control dynamics. Zero-forgetting was not universal across all TSN variants.
Routing Efficacy:
- TSN-ReplayKL (observation-based routing) performed poorly, worse than TSN-Core, indicating that observation similarity is insufficient for routing in continuous control.
- TSN-Affinity-L (Latent affinity) achieved the best results with zero forgetting and the lowest average gap from target returns (0.670), outperforming Cumulative Replay (0.729).
Trade-off: The best performance required creating multiple model copies (3 copies), highlighting a performance-capacity trade-off. However, the results confirm that RL-aware affinity is a reliable signal for routing in complex continuous domains where replay heuristics fail.

5. Significance and Conclusion

The paper demonstrates that architectural continual learning is a viable and often superior alternative to replay-based strategies in the offline RL setting.

Mitigation of Catastrophic Forgetting: By isolating task-specific parameters in sparse subnetworks, the method effectively prevents destructive interference.
Efficiency over Replay: TSN-Affinity avoids the memory overhead and distribution mismatch issues inherent in replay buffers.
Superior Routing: The study proves that control-relevant affinity scores (action and latent similarity) are far more effective for routing decisions than simple observation-based similarity or replay heuristics, especially in continuous control tasks.
Future Impact: This work paves the way for scalable, memory-efficient continual learning in safety-critical domains (like robotics) where re-collecting data or interacting with the environment is not feasible.

The authors conclude that similarity-guided architectural reuse is a robust strategy for CORL, offering a strong balance between retaining past knowledge and transferring it to new, related tasks.