Energy-Based Transfer for Reinforcement Learning

arXiv:2506.16590 · cs.LG, cs.AI · Submitted 2025-06-19 · Read on arXiv

Listen

Transcript

Introduction to the show: ident: AI Radio. Generated commentary on the latest Artificial Intelligence papers.

Tom: We're looking at the paper: "Energy-Based Transfer for Reinforcement Learning".

Jane: Here's the gist of the paper. The gist: We propose an energy-based transfer learning method that uses out-of-distribution detection to selectively issue guidance, enabling the teacher to intervene only in states within its training distribution.

Tom: Reinforcement learning algorithms often suffer from poor sample efficiency, making them challenging to apply in multi-task or continual learning settings.

Jane: Efficiency can be improved by transferring knowledge from a previously trained teacher policy to guide exploration in new but related tasks.

Tom: However, if the new task sufficiently differs from the teacher’s training task, the transferred guidance may be sub-optimal and bias exploration toward low-reward behaviors.

Paper summary: Tom: So we've got the idea that you can use energy scores to selectively issue guidance, meaning the teacher only steps in when it thinks it knows what it's doing. Jane That selective intervention is key because if you give a bad suggestion in a weird state, you just make things worse for learning. Lu The paper formally defines the free energy E(x; f) as minus T log X K i equals one e to f i over T, and they say an input x is out-of-distribution if its energy score is greater than a threshold tau >

Meng: So in plain terms, they treat states with high energy scores as in-distribution and therefore safe for the teacher to give an expert action suggestion. Tom That means the core mechanism checks if state s has an energy score phi(s; πT) greater than tau, and if so it suggests a known good action from the teacher policy πT >

Jane: And they introduce a loss function called Ltotal which combines the standard reinforcement learning loss LRL with an energy regularization term λ times Lenergy to make sure the teacher learns those energy scores well too. Lu They also mention that to make this detection better, they add a fixed set of representative out-of-distribution states into the teacher's training data, which helps separate in-distribution from out-of-distribution examples >

Tom: So they found that under on policy training, the log density of the teacher policy dπ(s) is proportional to that energy score phi(s) and also proportional to minus E(s), meaning high energy means high visitation density >

Jane: They tested this on two types of scenarios: single-task transfer and multi-task transfer, and they found that their method beats all the other methods in terms of how many samples you need to get a good policy result >

Lu: It’s interesting because they also looked at how changing the threshold tau affects things; initially increasing tau boosts performance by blocking bad advice but if tau gets too high, the teacher starts withholding guidance even in states it already knows well >

Meng: From an engineering standpoint I wonder about that threshold. How do you pick that tau value for a real system? Tom That’s what they address by saying they need to specify an energy threshold and it's primarily designed for settings with covariance rather than just label shift >

Jane: So the paper essentially shows how to make a teacher policy introspective, only helping when it's confident in its knowledge about the state you're in. Tom It’s a neat way to handle that transfer problem without having to train a whole new teacher for every new task >

Conclusion: Tom: So looking at "Energy-Based Transfer for Reinforcement Learning" by Deng et al., the authors are showing us how you can use energy scores as a reliable proxy for state familiarity in reinforcement learning. Jane They prove that when you only let the teacher intervene in states with high energy, you get better sample efficiency and higher returns across single-task and multi-task setups >

Lu: The main implication is that we don't need to blindly transfer knowledge; we can use a detection mechanism to filter out the bad advice when the new task drifts too far from what the teacher learned >

Meng: But they did note a limitation, which is that you have to specify an energy threshold tau, and they say it works best in settings with covariance instead of just label shift >

Tom: It’s important to remember that this paper shows the energy score is related to state visitation density under on-policy training, which provides the theoretical underpinning for why this detection method works >

Jane: So for someone who just listens, what it means is that you can build a system where an older version of an AI helps a newer version only when that newer version is in a situation the older one has seen before >

Zeyun Deng, Jasorsi Ghosh, Fiona Xie, Yuzhe Lu, Katia Sycara, Joseph Campbell

Purdue University · AWS AI Institute of Technology, Carnegie Mellon University

cs.LG, cs.AI

Submitted: 2025-06-19

Updated: 2026-10-20

Comments: Accepted to CoLLAs 2026

License: http://creativecommons.org/licenses/by/4.0/

Importance score: 78/100

The gist: The gist: We propose an energy-based transfer learning method that uses out-of-distribution detection to selectively issue guidance, enabling the teacher to intervene only in states within its

Key concepts

Energy-Based Out-of-Distribution Detection: This technique uses a mathematical function called free energy, derived from a neural network's output, to measure how 'out of distribution' (OOD) a state is. States with high energy scores are considered in-distribution and eligible for guidance, while states with low scores are OOD.
Energy Score ($\phi(s; \pi_T)$): The energy score is defined as the negative free energy of a state $s$ under the teacher policy $\pi_T$. It acts as a proxy for how likely that specific state is to have occurred during the teacher's original training. High scores suggest familiarity, making them good candidates for guidance.
Energy-Based Transfer Learning: This mechanism uses the energy score to decide whether to apply guidance. If the state's energy score is above a threshold ($\tau$), the student follows the teacher's action; otherwise, it explores independently. This ensures teacher intervention is selective and only occurs in states where its knowledge is expected to be valid.

Terminology

Reinforcement Learning — 2× in this paper · explained in 4 episodes
Sample Efficiency — 2× in this paper · explained in 7 episodes
Continual Learning — 1× in this paper · explained in 3 episodes
Transfer Learning — 1× in this paper · explained in 16 episodes

Summary

The gist: We propose an energy-based transfer learning method that uses out-of-distribution detection to selectively issue guidance, enabling the teacher to intervene only in states within its training distribution.

Introduction and Motivation

Reinforcement learning algorithms often suffer from poor sample efficiency, making them challenging to apply in multi-task or continual learning settings. Efficiency can be improved by transferring knowledge from a previously trained teacher policy to guide exploration in new but related tasks. However, if the new task sufficiently differs from the teacher’s training task, the transferred guidance may be sub-optimal and bias exploration toward low-reward behaviors. In this paper, we introduce an introspective transfer learning method that selectively guides exploration only when the teacher’s knowledge is likely to be helpful. Our approach, energy-based transfer learning (EBTL), is based on the insight that guidance should only be issued when the student visits states that lie within the teacher’s training distribution.

Energy-Based Out-of-Distribution Detection

The method leverages concepts from energy-based learning and out-of-distribution detection to compute energy scores over states visited by the student during training, treating high-energy states as in-distribution and therefore eligible for guidance. Formally, given an input x ∈ R D and a neural network f(x) ∈ R K with logits f1(x),..., fK(x), we define the free energy for x as E(x; f) = −T logX K i=1 e(fi(x)/T). An input is considered to be OOD if E(x; f) > τ for an energy threshold τ and in-distribution (ID) otherwise. The score function is defined as ϕ(s; πT) = −E(s; πT), where this serves as a proxy for how likely the state is to belong to the teacher’s training distribution p(x). Proposition 4.1 states that under on-policy training, log dπ(s) ∝ ϕ(s) and log dπ(s) ∝ −E(s).

Energy-Based Transfer Learning Mechanism

The action selection rule is defined as a conditional logic based on the energy score: a = πT (· s), if ϕ(s; πT) ≥ τ, and aS ∼ πS(· s), if ϕ(s; πT) < τ. The teacher checks if each state is ID or OOD by comparing the state’s energy score to a pre-defined energy threshold. To improve separability, we adopt an energy-based loss that augments the teacher’s training with a fixed set of representative OOD states. The overall teacher loss is defined as Ltotal = LRL + λ · Lenergy where λ ∈ R+ controls the weight of the energy regularization.

Experimental Results and Findings

The method was evaluated in two settings: single-task transfer and multi-task transfer. In both settings, EBTL consistently outperforms all baselines in terms of sample efficiency and final policy return across all scenarios. For the Alternating Goal Room, when the energy threshold q ≥ 0.1, EBTL rarely issues guidance in unfamiliar states, leading to significant improvements in transfer performance. The results show that increasing q initially boosts transfer performance by suppressing harmful advice but when q becomes too large, the teacher begins to withhold guidance even in familiar situations.

Conclusion and Limitations

EBTL consistently outperforms all baselines in both sample efficiency and final policy return across all scenarios. The paper concludes that EBTL leverages energy scores as a proxy for familiarity, issuing advice only in states likely to lie within the teacher’s training distribution. A limitation noted is that EBTL requires specifying an energy threshold and is primarily designed for settings with covariance rather than label shift. The method demonstrates robustness to layout complexity in multi-task settings, yielding high returns and low variance across seeds.

Appendix Details

In the GridWorld experiments, action masking is applied to dynamically restrict the agent’s action space based on its immediate environment. The teacher policy trained with energy loss uses min and mout set to 10 and 15 respectively. For the Overcooked-AI experiments, shaped rewards are introduced to facilitate the training process. The student training horizon varies depending on the setup and source-target configuration. The sensitivity evaluation protocol shows no statistically significant variation in the separation between false OOD and true OOD energy distributions across different regularization settings. The Kruskal–Wallis tests yield p-values above 0.1 for all four divergence metrics, indicating that the teacher’s ability to distinguish between states where guidance should or should not be issued is stable across regularization settings. The final results show EBTL maintains positive transfer under increasing covariate shift in multi-task settings.

References

[1] Marcin Andrychowicz, Filip Wolski, Alex Ray, Jonas Schneider, Rachel Fong, Peter Welinder, Bob McGrew, Josh Tobin, OpenAI Pieter Abbeel, and Wojciech Zaremba. Hindsight experience replay.

[2] Kai Arulkumaran, Marc Peter Deisenroth, Miles Brundage, and Anil Anthony Bharath. Deep reinforcement learning: A brief survey.

[3] Albert Bou, Matteo Bettini, Sebastian Dittert, Vikash Kumar, Shagun Sodhani, Xiaomeng Yang, Gianni De Fabritiis, and Vincent Moens. Torchrl: A data-driven decision-making library for pytorch.

[4] Tim Brys, Anna Harutyunyan, Matthew E Taylor, and Ann Nowé. Policy transfer using reward shaping.

[5] Joseph Campbell, Yue Guo, Fiona Xie, Simon Stepputtis, and Katia Sycara. Introspective action advising for interpretable transfer learning.

[6] Micah Carroll, Rohin Shah, Mark K Ho, Tom Griffiths, Sanjit Seshia, Pieter Abbeel, and Anca Dragan. On the utility of learning about humans for human-ai coordination.

[7] Maxime Chevalier-Boisvert, Bolun Dai, Mark Towers, Rodrigo Perez-Vicente, Lucas Willems, Salem Lahlou, Suman Pal, Pablo Samuel Castro, and Jordan Terry. Minigrid & miniworld: Modular & customizable reinforcement learning environments for goal-oriented tasks.

[8] Yue Guo, Joseph Campbell, Simon Stepputtis, Ruiyu Li, Dana Hughes, Fei Fang, and Katia Sycara. Explainable action advising for multi-agent reinforcement learning.

[9] Yue Guo, Xijia Zhang, Simon Stepputtis, Joseph Campbell, and Katia P. Sycara. Adaptive action advising with different rewards.

[10] Tuomas Haarnoja, Haoran Tang, Pieter Abbeel, and Sergey Levine. Reinforcement learning with deep energy-based policies.

[11] Anna Harutyunyan, Sam Devlin, Peter Vrancx, and Ann Nowé.

Improvements for AI systems

Energy based transfer learning enables selective guidance by computing energy scores over states visited by the student during training, treating high-energy states as in-distribution and therefore eligible for guidance. This mechanism allows the teacher to act only when it is sufficiently 'familiar' with the current context, leading to more efficient training not by issuing more guidance but by issuing correct guidance.
The energy score acts as a proxy for how likely a state is to belong to the teacher’s training distribution, as proven by Proposition 4.1, which states that the log of the visitation density is proportional to the score function phi(s) = -E(s). This allows the system to theoretically justify why guidance should be issued based on state familiarity.
The system can improve sample efficiency and performance in multi-task learning by selectively guiding exploration only when the student’s state lies within the teacher’s training distribution. This is achieved via Equation 2, which restricts teacher intervention to states where the energy score is greater than a threshold tau, enabling positive transfer even between tasks with high degrees of covariate shift.
The model can be trained to distinguish between false and true out-of-distribution states by augmenting teacher training with an energy based loss that penalizes ID states with low scores and OOD states with high scores. This regularization helps improve separability, though the paper notes that no statistically significant variation in the separation between false OOD and true OOD energy distributions across different (min, mout) configurations was observed.
The system can be implemented in complex environments like Overcooked where it maintains positive transfer under increasing covariate shift by restricting guidance to states associated with recipes the teacher has encountered during training. This robustness allows for learning new task variants even when they introduce both recipe and layout shift.

Abstract

Reinforcement learning remains sample inefficient, particularly when agents must adapt over time to related but distinct tasks, as in continual learning settings. Efficiency can be improved by transferring knowledge from a pretrained teacher policy to guide exploration in new but related tasks. However, if the new task sufficiently differs from the teacher's training task, the transferred guidance may be sub-optimal and bias exploration toward low-reward behaviors. We propose an energy-based transfer learning method that uses out-of-distribution detection to selectively issue guidance, enabling the teacher to intervene only in states within its training distribution. We theoretically link energy scores to state-visitation patterns under on-policy reinforcement learning. Empirically, our method improves sample efficiency and performance across both single-task and multi-task settings.