Knowledge Distillation for Efficient Transformer-Based Reinforcement Learning in Hardware-Constrained Energy Management Systems

Pascal Henrich; Jonas Sievers; Maximilian Beichter; Thomas Blank; Ralf Mikut; Veit Hagenmeyer

cs.LGMar 27, 2026

Knowledge Distillation for Efficient Transformer-Based Reinforcement Learning in Hardware-Constrained Energy Management Systems

Authors: Pascal Henrich, Jonas Sievers, Maximilian Beichter, Thomas Blank, Ralf Mikut, Veit Hagenmeyer

Organizations: Karlsruhe Institute of Technology, Karlsruhe, Germany

Abstract

Transformer-based reinforcement learning has emerged as a strong candidate for sequential control in residential energy management. In particular, the Decision Transformer can learn effective battery dispatch policies from historical data, thereby increasing photovoltaic self-consumption and reducing electricity costs. However, transformer models are typically too computationally demanding for deployment on resource-constrained residential controllers, where memory and latency constraints are critical. This paper investigates knowledge distillation to transfer the decision-making behaviour of high-capacity Decision Transformer policies to compact models that are more suitable for embedded deployment. Using the Ausgrid dataset, we train teacher models in an offline sequence-based Decision Transformer framework on heterogeneous multi-building data. We then distil smaller student models by matching the teachers' actions, thereby preserving control quality while reducing model size. Across a broad set of teacher-student configurations, distillation largely preserves control performance and even yields small improvements of up to 1%, while reducing the parameter count by up to 96%, the inference memory by up to 90%, and the inference time by up to 63%. Beyond these compression effects, comparable cost improvements are also observed when distilling into a student model of identical architectural capacity. Overall, our results show that knowledge distillation makes Decision Transformer control more applicable for residential energy management on resource-limited hardware.

Figures & tables

Figure 1. Graphical Abstract

Figure 2. Markov Decision Process for Reinforcement Learning in a Battery Scheduling Environment.

Figure 3. Decision Transformer in a Battery Scheduling Environment (adapted from ( Chen et al., 2021 ) ).

Figure 4. Response-Based Knowledge Distillation for Decision Transformer in a Battery Scheduling Environment.

Figure 5. Aggregated Daily Mean Load for Building 13 with a Typical Residential Load Pattern.

Name	# Layers	# Heads	Model Dimension
Tiny	1	1	64
Mini	2	2	128
Small	3	1	128
Medium	8	2	256
Large	12	4	512

Table 1. Hyperparameter of Different Decision Transformer Model Sizes.

Figure 6. Mean Cost Across All 20 Buildings for Different Approaches. DT Denotes the Medium Decision Transformer, while KD Refers to Knowledge Distillation From Medium Teacher to Small Student.

Figure 7. Mean Cost Across All 20 Buildings for Different Decision Transformer Sizes.

Figure 8. Mean Cost Improvement Across All 20 Buildings for Different Knowledge Distillation Sizes.

Figure 9. Best Performing Learning-Based Approach for Each of the 20 Buildings. DT Denotes the Medium Decision Transformer, while KD Refers to Knowledge Distillation From Medium Teacher to Small Student.

	Total Parameters	Total Memory [ $\text{\,}\mathrm{MB}$ ]	Time [ms]
Tiny	$387\,713$	$7.09$	$3.56$
Mini	$1\,498\,241$	$17.44$	$4.80$
Small	$2\,157\,569$	$25.45$	$5.99$
Medium	$9\,831\,681$	$70.56$	$9.66$
Large	$51\,166\,721$	$219.83$	$16.83$

Table 2. Model Size Comparison in Terms of Total Parameters, Memory Footprint, and Inference Time for Different Decision Transformer Variants.

Appendix figures & tables6 assets

Supplementary material from the paper’s appendix.

Appendix

Hyperparameter	Value
Batch size	32
Forecast Horizon	24
Context Length	96
Optimiser DT	Adam
Optimiser DT Learning Rate	1e-4
Optimiser DT Weight Decay	1e-4

Appendix

Table 3. Hyperparameters for Decision Transformer & Knowledge Distillation.

Figure 10. Best Performing Learning-Based Approach for Each of the 20 Buildings. DT Denotes the Medium Decision Transformer, while KD Refers to Knowledge Distillation From Medium Teacher to Tiny Student.

Bld. ID	MILP	WO Battery	DDPG	DT Tiny	DT Mini	DT Small	DT Medium	DT Large
13	181.72	230.65	209.88 ± 4.8	210.79 ± 1.49	210.6 ± 1.62	210.68 ± 1.8	212.52 ± 0.66	210.29 ± 2.87
20	357.44	392.74	384.01 ± 3.11	380.62 ± 1.22	381.22 ± 0.88	381.13 ± 1.0	380.34 ± 1.5	379.88 ± 0.4
33	132.88	167.11	152.47 ± 2.03	151.81 ± 1.33	151.43 ± 0.73	152.18 ± 0.63	151.25 ± 1.33	152.08 ± 1.54
35	314.47	343.74	334.37 ± 2.85	330.8 ± 1.95	330.45 ± 1.53	330.34 ± 1.12	329.38 ± 2.11	329.33 ± 0.97
74	205.15	289.89	256.58 ± 1.33	258.81 ± 4.53	255.26 ± 3.47	255.7 ± 3.01	254.9 ± 3.78	266.91 ± 6.73
75	70.97	206.03	170.22 ± 8.77	160.28 ± 6.81	156.06 ± 6.77	157.26 ± 4.14	154.93 ± 6.09	169.38 ± 9.64

Appendix

Table 4. Cost in € Across All 20 Buildings for the Three Benchmarks and Five DT Variants, Averaged Over Three Random Seeds. The Final Row Reports the Mean Cost for Each Approach.

Bld. ID	Small Teacher	Tiny Student	Mini Student	Small Student
13	210.68 ± 1.8	210.51 ± 1.28	210.35 ± 2.49	209.93 ± 1.87
20	381.13 ± 1.0	380.65 ± 1.0	380.21 ± 0.42	380.19 ± 0.16
33	152.18 ± 0.63	151.84 ± 1.01	151.22 ± 0.4	151.36 ± 0.8
35	330.34 ± 1.12	330.25 ± 0.71	329.5 ± 0.66	329.49 ± 0.58
74	255.7 ± 3.01	254.8 ± 1.25	253.84 ± 2.09	253.56 ± 1.83
75	157.26 ± 4.14	159.02 ± 6.36	155.81 ± 6.13	155.21 ± 5.57

Appendix

Table 5. Cost in € Across All 20 Buildings for the Small Teacher and the Corresponding Student Models, Averaged Over Three Random Seeds. The Final Row Reports the Mean Cost for Each Approach.

Bld. ID	Medium Teacher	Tiny Student	Mini Student	Small Student
13	212.52 ± 0.66	210.29 ± 1.17	211.02 ± 1.08	211.42 ± 0.45
20	380.34 ± 1.5	379.61 ± 0.04	379.25 ± 0.46	379.2 ± 0.11
33	151.25 ± 1.33	150.45 ± 0.44	150.56 ± 0.22	150.61 ± 0.54
35	329.38 ± 2.11	329.07 ± 1.43	328.85 ± 1.46	328.6 ± 1.32
74	254.9 ± 3.78	254.12 ± 2.44	254.99 ± 4.33	253.49 ± 2.55
75	154.93 ± 6.09	153.1 ± 5.81	152.8 ± 5.89	152.19 ± 6.45

Appendix

Table 6. Cost in € Across All 20 Buildings for the Medium Teacher and the Corresponding Student Models, Averaged Over Three Random Seeds. The Final Row Reports the Mean Cost for Each Approach.

Bld. ID	Large Teacher	Tiny Student	Mini Student	Small Student
13	210.29 ± 2.87	209.79 ± 0.4	211.61 ± 0.41	210.36 ± 2.16
20	379.88 ± 0.4	379.91 ± 0.96	379.52 ± 0.23	379.4 ± 0.36
33	152.08 ± 1.54	150.59 ± 1.01	150.65 ± 0.64	150.42 ± 0.63
35	329.33 ± 0.97	329.47 ± 0.49	329.37 ± 0.66	329.22 ± 0.71
74	266.91 ± 6.73	262.12 ± 8.92	262.58 ± 9.37	262.96 ± 8.49
75	169.38 ± 9.64	164.06 ± 11.6	162.35 ± 14.59	162.47 ± 15.79

Appendix

Table 7. Cost in € Across All 20 Buildings for the Large Teacher and the Corresponding Student Models, Averaged Over Three Random Seeds. The Final Row Reports the Mean Cost for Each Approach.

Explore similar work

Interpretable Policy Distillation for Power Grid Topology Control
May 30, 2026Aleksandra Dmitruka, Karlis FreivaldsRL Control Reinforcement Learning
University of Latvia, Faculty of Exact Sciences and Technology, Riga, Latvia
Adaptive Model Compression (AMC): Saliency-Driven Resource Allocation for Ultra-Low-Power Transformer Inference
Jul 11, 2026Jiayin Hu, Kai Yuan, Vanessa Hu +3Efficient Transformer Inference Model Compression
Apple USA
Topology-Aware Graph Reinforcement Learning for Energy Storage Systems Optimal Dispatch in Distribution Networks
Mar 27, 2026Shuyi Gao, Stavros Orfanoudakis, Shengren Hou +2Power Systems RL Control
Department of Electrical Sustainable Energy, Delft University of Technology, Mekelweg 4, Delft, 2628 CD, The Netherlands · Beijing Energy Quant Ltd., Haidian, Beijing, 53100, China

May 30, 2026J/K move · Enter open · S save

Interpretable Policy Distillation for Power Grid Topology Control

Aleksandra Dmitruka, Karlis Freivalds

University of Latvia, Faculty of Exact Sciences and Technology, Riga, Latvia

RL Control Reinforcement Learning Decision Tree Learning Interpretability in RL Policy Distillation Power Systems

Open paper AskPDF arXiv

Deep reinforcement learning (RL) offers a promising route to real-time power grid operation, yet large neural policies are costly to evaluate, hard to deploy on constrained hardware, and opaque to operators. We ask whether a Proximal Policy Optimization (PPO) agent for grid topology control can be compressed into compact tree-based surrogates without losing operational performance. A PPO teacher is trained on Grid2Op's standard 14-bus environment with a stability-oriented reward, using stress-focused data collection on critical, high-loading states. The policy is then distilled into a decision tree and a random forest. Across held-out validation episodes, both surrogates exceed the teacher in mean reward and survival length at a fraction of the inference cost. The decision tree shows high exact-action agreement with the PPO argmax and near-complete agreement within its top-ranked actions, while remaining small enough to be inspected directly. Feature-importance analysis reveals a representational shift: the PPO policy relies mainly on line-loading signals, while the distilled tree is driven primarily by bus-topology variables. These results suggest that stress-focused distillation can convert a black-box neural controller into a lightweight, auditable rule-like surrogate suited for real-time deployment, while also surfacing risks tied to deterministic actions and topology-specific generalization.

Abstract

Explore similar work

Interpretable Policy Distillation for Power Grid Topology Control

Adaptive Model Compression (AMC): Saliency-Driven Resource Allocation for Ultra-Low-Power Transformer Inference

Topology-Aware Graph Reinforcement Learning for Energy Storage Systems Optimal Dispatch in Distribution Networks