cs.LGJun 19, 2025
SaveEnergy-Based Transfer for Reinforcement Learning
Organizations: ♠Purdue University · ♢AWS AI · ♣Carnegie Mellon University
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.
Figures & tables
Figure 1: Overview of energy-based transfer learning . As the student interacts with the environment, the teacher: 1) checks if each state is in-distribution or out-of-distribution by comparing the state’s energy score to a pre-defined energy threshold ; 2) If the state exceeds the energy threshold , then it is considered in-distribution for the teacher and an expert action is suggested to the student.
Figure 2: Environments used for empirical experiments. See Section 5 for detailed descriptions.
Figure 3: Minigrid (10 seeds). (a) Empirical energy score distributions with respect to the teacher policy. The source task (blue) shows the teacher’s training distribution. The target task (orange + purple), measured during transfer, is bimodal: one mode overlaps with the source (shared sub-task, in-distribution), while the other does not (non-shared sub-task, out-of-distribution). (b) Guidance quality: fraction of transferable target states correctly identified as ID and advised (orange), and fraction of non-transferable target states incorrectly identified as ID and advised (purple), with the decay schedule disabled. An ideal gating signal yields an orange bar substantially higher than the purple bar. (c) Relative transfer (% vs. scratch) across and decay schedules (50%: guidance reaches 0 at mid-training). denotes always advising (AA baseline). (d) Learning curves: EBTL vs. baselines. Guidance ends at mid-training with .
Figure 4: Lang-Minigrid . (10 seeds). Guidance ends at mid-training with . Task variation is specified through language instructions: in Alternating-Goal, the instruction indicates the target corner (e.g., go to the top left corner ), while in Unlocked-to-Locked, it specifies sequential subgoals (e.g., pick up the key, then go to the goal ).
Figure 5: (10 seeds). Learning curves with (solid) and without (dashed) energy regularization in Minigrid environments.
Figure 6: OOD score distributions across detection methods in Minigrid. Each panel shows the score distributions of the source task (blue) and target task (orange: shared sub-task, purple: non-shared sub-task) under the teacher policy. A good OOD detector should separate blue and purple while keeping blue and orange overlapping. (a) MaxP, (b) Entropy, and (c) ODIN, requiring OOD samples for hyperparameter calibration. (d) (Left) The energy score without regularization (no OOD samples) already matches ODIN, and (Right) with regularization (OOD samples used during teacher training) achieves substantially sharper separation.
| Optimality | EBTL (0.3) | EBTL (0.5) | EBTL (0.7) | AA | JSRL | KSRL | Fine-tuning |
|---|---|---|---|---|---|---|---|
| 96% | |||||||
| 70% | |||||||
| 45% | |||||||
| 15% |
Table 1: Robustness to suboptimal teachers . Relative transfer improvement (%) over training from scratch (mean 95% CI, 10 seeds) for Alternating Goal. Guidance ends at mid-training.
Figure 7: Overcooked (5 seeds). (Cols 1 & 3) Energy scores: source (blue) vs. target (orange + purple). Target is bimodal: one mode overlaps source (ID), one does not (OOD). (Cols 2 & 4) Learning curves: EBTL vs. baselines ( for Recipe Shift, for Recipe + Layout Shift).
Appendix figures & tables20 assets
Supplementary material from the paper’s appendix.
Appendix
| Setup | Recipe (2 3) | Recipe (1 3) | Recipe + Layout (2 2) |
|---|---|---|---|
| Simple Room | 19,008,000 | 9,004,800 | 12,000,000 |
| Ring Room | 2,400,000 | 10,003,200 | 18,000,000 |
Table 2: Selected teacher checkpoints (in environment steps) for each Overcooked setup and source-target configuration.
| Setup | Recipe (2 3) | Recipe (1 3) | Recipe + Layout (2 2) |
|---|---|---|---|
| Simple Room | 20M | 20M | 12M |
| Ring Room | 35M | 35M | 20M |
Table 3: Transfer horizons (in millions of environment steps) used for student training in each Overcooked setup and configuration. Each experiment is run with 5 random seeds.
| Hyperparameter | Value |
| Learning rate | 0.0005 |
| Discount factor ( ) | 0.9 |
| GAE lambda ( ) | 0.8 |
| Policy clip parameter | 0.2 |
| Value function clip parameter | 10.0 |
| Value loss coefficient | 0.5 |
Table 4: Hyperparameters used for all GridWorld experiments.
Figure 8: 10 seeds. JSRL hyperparameter tuning in MiniGrid.
Figure 9: 10 seeds. KSRL hyperparameter tuning in MiniGrid.
Figure 10: Results (3 seeds). In each environment, the left column shows the empirical density of the teacher’s state energy . Blue traces the source rollouts used to train the teacher. Orange and green are target rollouts during transfer and typically form a bimodal pattern: one component overlaps the source shared task (in distribution), while the other occupies a separate region (non-shared task, out of distribution). is set to 0.3 in both settings.
Figure 11: Actor-Critic architectures used in our experiments. (a) MiniGrid. (b) Overcooked.
| Metric | ANOVA p-value | Kruskal–Wallis p-value |
|---|---|---|
| Jensen–Shannon | 0.1138 | 0.1592 |
| Kullback–Leibler | 0.2457 | 0.1799 |
| Total Variation | 0.1728 | 0.2322 |
| Hellinger Distance | 0.1247 | 0.1592 |
Table 7: Statistical test results (p-values) for divergence between False OOD and True OOD energy distributions across different settings.
Table 8: Choice of ID states for setting the quantile . “Post-train” derives from evaluation rollouts of the converged teacher (no exploration); “In-train” derives from exploration-time states during teacher training. Values are student transfer performance relative to training from scratch (mean 95% CI)
Figure 12: 10 seeds. Effect of replacing rollout-based OOD batches with randomly sampled OOD batches. advise in all states (AA Baseline).
| Teacher | Without energy loss | With energy loss |
|---|---|---|
| Alternating Room | 100,000 | 60,000 |
| Unlocked-to-Locked | 260,000 | 220,000 |
Table 9: Training steps to convergence (mean over 10 seeds; lower is better).
Figure 13: 10 seeds. Comparison of five decay schedules. advise in all states (AA Baseline).
Figure 14: Heatmaps of average energy quantiles under the teacher for Alternating Goal. Left: source (goal in Room 1). Right: target (goal in Room 3). Higher quantiles indicate greater teacher familiarity.
Figure 15: 10 seeds. Advice schedules where advice ends at (a) 25%, (b) 50%, and (c) 75% of training in the Unlocked-to-Locked environment. advise in all states (AA Baseline).
Figure 16: Learning curves across (10 seeds). EBTL for vs. baselines; near-identical across the range. Guidance ends at mid-training.
Figure 17: IAA reproduction and decay-schedule comparison. 10 seeds. Top: exponential decay (as in the original IAA paper). Bottom: linear decay (the EBTL default). Under exponential decay we reproduce the IAA result, but standard action advising without filtering also exceeds the baseline, suggesting the schedule itself contributes substantially to the gain. Under linear decay, only EBTL remains clearly above the baseline. EBTL is the best performing method under both schedules.
Figure 18: IAA with extended burn-in periods. 10 seeds. Linear Decay. Learning curves for IAA across burn-in start points (25% and 50% of training), advice decay endpoints (50%, 75%, and 100% of training), and thresholds ( ).
Figure 19: IAA value-gap separation in the Unlocked-to-Locked setting , shown across (one per row); left column freezes the convolutional backbone, right column fine-tunes the full network. Each panel tracks three groups of states over the course of the student’s learning: ID (in-distribution, collected during teacher training), Transferable OOD (states encountered during the student’s learning where the teacher should issue advice), and Non-transferable OOD (states where the teacher should not ). A reliable signal should place the ID and Transferable-OOD curves above the Non-transferable-OOD curve, indicating that good advice is issued more often than bad advice.
Figure 20: IAA value-gap separation in the Alternating-Goal setting , across (one per row); left column freezes the convolutional backbone, right column fine-tunes the full network. Curves and the expected ordering are as in Figure 19 : a reliable signal keeps the ID and Transferable-OOD curves above the Non-transferable-OOD curve.